From 89bbb94ad5806b4429fd8151917d3a3c081dde14 Mon Sep 17 00:00:00 2001
From: Eric Zarowny <ezarowny@gmail.com>
Date: Thu, 4 Aug 2016 14:43:37 -0700
Subject: [PATCH] better enum support + fix for enum34 IntEnum quirk

---
 rollbar/lib/transforms/__init__.py          | 13 +++++--
 rollbar/test/test_serializable_transform.py | 38 ++++++++++++++++++++-
 setup.py                                    | 20 +++++++----
 3 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/rollbar/lib/transforms/__init__.py b/rollbar/lib/transforms/__init__.py
index a7d4e4e1..4bd9876e 100644
--- a/rollbar/lib/transforms/__init__.py
+++ b/rollbar/lib/transforms/__init__.py
@@ -79,12 +79,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, integer_types + (float,)):
+            if type(o) not in integer_types + (float, ):
+                return do_transform('custom', o, key=key)
+            else:
+                return do_transform('number', o, key=key)
+
         return do_transform('custom', o, key=key)
 
     handlers = {
diff --git a/rollbar/test/test_serializable_transform.py b/rollbar/test/test_serializable_transform.py
index 782534dd..8ff63e30 100644
--- a/rollbar/test/test_serializable_transform.py
+++ b/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
diff --git a/setup.py b/setup.py
index 4c4954f3..63dc5f21 100644
--- a/setup.py
+++ b/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,6 +17,17 @@
     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(),
@@ -65,11 +77,5 @@
         'requests',
         'six',
         ],
-    tests_require=[
-        'mock',
-        'webob',
-        'blinker',
-        'unittest2'
-        ],
+    tests_require=tests_require,
     )
-