Merge pull request #121 from rollbar/better-enum-support

better enum support + fix for enum34 IntEnum quirk

rollbar/lib/__init__.py

@@ -1,4 +1,5 @@
 import base64
+import collections
 import copy
 import os
 import sys
@@ -11,7 +12,9 @@
 
 binary_type = six.binary_type
 integer_types = six.integer_types
+number_types = integer_types + (float, )
 string_types = six.string_types
+sequence_types = (collections.Mapping, list, tuple, set, collections.deque)
 
 urlparse = urllib.parse.urlparse
 urlsplit = urllib.parse.urlsplit
@@ -44,25 +47,23 @@ def text(val):
 
         return repr(val)
 
-
     _map = map
+
     def map(*args):
         return _map(*args)
 
-
     def force_lower(val):
         return str(val).lower()
 
 else:
     def text(val):
         return str(val)
 
-
     _map = map
+
     def map(*args):
         return list(_map(*args))
 
-
     def force_lower(val):
         try:
             return val.lower()

rollbar/lib/transforms/__init__.py

@@ -1,4 +1,6 @@
-from rollbar.lib import python_major_version, binary_type, string_types, integer_types, traverse
+from rollbar.lib import (
+    python_major_version, binary_type, string_types, integer_types,
+    number_types, traverse)
 
 _ALLOWED_CIRCULAR_REFERENCE_TYPES = [binary_type, bool, type(None)]
 
@@ -79,12 +81,19 @@ def string_handler(s, key=None):
                 return do_transform('unicode', s, key=key)
 
     def default_handler(o, key=None):
-        if isinstance(o, integer_types + (float,)):
-            return do_transform('number', o, key=key)
-
         if isinstance(o, bool):
             return do_transform('boolean', o, key=key)
 
+        # There is a quirk in the current version (1.1.6) of the enum
+        # backport enum34 which causes it to not have the same
+        # behavior as Python 3.4+. One way to identify IntEnums is that
+        # they are instances of numbers but not number types.
+        if isinstance(o, number_types):
+            if type(o) not in number_types:
+                return do_transform('custom', o, key=key)
+            else:
+                return do_transform('number', o, key=key)
+
         return do_transform('custom', o, key=key)
 
     handlers = {

rollbar/lib/transforms/serializable.py

@@ -1,7 +1,9 @@
 import math
 
 from rollbar.lib import binary_type, string_types
-from rollbar.lib import circular_reference_label, float_infinity_label, float_nan_label, undecodable_object_label, unencodable_object_label
+from rollbar.lib import (
+    circular_reference_label, float_infinity_label, float_nan_label,
+    undecodable_object_label, unencodable_object_label)
 from rollbar.lib import iteritems, python_major_version, text
 
 from rollbar.lib.transforms import Transform
@@ -106,7 +108,8 @@ def transform_custom(self, o, key=None):
             try:
                 exc_str = str(e)
             except Exception as e2:
-                exc_str = '[%s while calling str(%s)]' % (e2.__class__.__name__, e.__class__.__name__)
+                exc_str = '[%s while calling str(%s)]' % (
+                    e2.__class__.__name__, e.__class__.__name__)
             return '<%s in %s.__repr__: %s>' % (
                 e.__class__.__name__, o.__class__.__name__, exc_str)
 

rollbar/lib/transforms/shortener.py

@@ -1,8 +1,9 @@
 from array import array
 import collections
-import math
 
-from rollbar.lib import integer_types, iteritems, key_in, reprlib, string_types, text
+from rollbar.lib import (
+    integer_types, iteritems, key_in, number_types, reprlib, sequence_types,
+    string_types, text)
 from rollbar.lib.transforms import Transform
 
 
@@ -68,10 +69,10 @@ def _shorten_other(self, obj):
     def _shorten(self, val):
         max_size = self._get_max_size(val)
 
-        if isinstance(val, (string_types, collections.Mapping, list, tuple, set, collections.deque)):
+        if isinstance(val, (string_types, sequence_types)):
             return self._shorten_sequence(val, max_size)
 
-        if isinstance(val, integer_types):
+        if isinstance(val, number_types):
             return self._shorten_basic(val, self._repr.maxlong)
 
         return self._shorten_other(val)

rollbar/test/test_serializable_transform.py

@@ -1,7 +1,7 @@
 import collections
 import base64
 import copy
-import math
+import enum
 
 from rollbar.lib import transforms, python_major_version
 from rollbar.lib.transforms.serializable import SerializableTransform
@@ -56,6 +56,42 @@ def test_simple_dict(self):
         expected = copy.deepcopy(start)
         self._assertSerialized(start, expected)
 
+    def test_enum(self):
+        class EnumTest(enum.Enum):
+            one = 1
+            two = 2
+
+        start = EnumTest.one
+        expected = "<enum 'EnumTest'>"
+        self._assertSerialized(start, expected)
+
+    def test_enum_no_safe_repr(self):
+        class EnumTest(enum.Enum):
+            one = 1
+            two = 2
+
+        start = EnumTest.one
+        expected = '<EnumTest.one: 1>'
+        self._assertSerialized(start, expected, safe_repr=False)
+
+    def test_int_enum(self):
+        class EnumTest(enum.IntEnum):
+            one = 1
+            two = 2
+
+        start = EnumTest.one
+        expected = "<enum 'EnumTest'>"
+        self._assertSerialized(start, expected)
+
+    def test_int_enum_no_safe_repr(self):
+        class EnumTest(enum.IntEnum):
+            one = 1
+            two = 2
+
+        start = EnumTest.one
+        expected = '<EnumTest.one: 1>'
+        self._assertSerialized(start, expected, safe_repr=False)
+
     def test_encode_dict_with_invalid_utf8(self):
         start = {
             'invalid': invalid

setup.py

@@ -1,5 +1,6 @@
 import re
 import os.path
+import sys
 from setuptools import setup, find_packages
 
 HERE = os.path.abspath(os.path.dirname(__file__))
@@ -16,11 +17,22 @@
     INIT_DATA = fd.read()
     VERSION = re.search(r"^__version__ = ['\"]([^'\"]+)['\"]", INIT_DATA, re.MULTILINE).group(1)
 
+tests_require = [
+    'mock',
+    'webob',
+    'blinker',
+    'unittest2'
+    ]
+
+version = sys.version_info
+if version[0] == 2 or (version[0] == 3 and version[1] < 4):
+    tests_require.append('enum34')
+
 setup(
     name='rollbar',
     packages=find_packages(),
     version=VERSION,
-    entry_points= {
+    entry_points={
         'paste.filter_app_factory': [
             'pyramid=rollbar.contrib.pyramid:create_rollbar_middleware'
         ],
@@ -65,11 +77,5 @@
         'requests',
         'six',
         ],
-    tests_require=[
-        'mock',
-        'webob',
-        'blinker',
-        'unittest2'
-        ],
+    tests_require=tests_require,
     )
-