Repository: rollbar/pyrollbar

Issue:
Ability to transform payload
Implement a transform function, similar to rollbar.js

Discussion:

Patch:
diff --git a/rollbar/__init__.py b/rollbar/__init__.py
index 2a3a62d..d26f6fd 100644
--- a/rollbar/__init__.py
+++ b/rollbar/__init__.py
@@ -323,6 +323,7 @@ SETTINGS = {
     'request_pool_maxsize': None,
     'request_max_retries': None,
     'batch_transforms': False,
+    'custom_transforms': [],
 }
 
 _CURRENT_LAMBDA_CONTEXT = None
@@ -426,12 +427,23 @@ def init(access_token, environment='production', scrub_fields=None, url_fields=N
                                    keys=shortener_keys,
                                    **SETTINGS['locals']['sizes'])
     _transforms = [
-        shortener,
-        ScrubRedactTransform(),
-        _serialize_transform,
-        ScrubTransform(suffixes=[(field,) for field in SETTINGS['scrub_fields']], redact_char='*'),
-        ScrubUrlTransform(suffixes=[(field,) for field in SETTINGS['url_fields']], params_to_scrub=SETTINGS['scrub_fields'])
+        shortener,  # priority: 10
+        ScrubRedactTransform(),  # priority: 20
+        _serialize_transform,  # priority: 30
+        ScrubTransform(suffixes=[(field,) for field in SETTINGS['scrub_fields']], redact_char='*'),  # priority: 40
+        ScrubUrlTransform(
+            suffixes=[(field,) for field in SETTINGS['url_fields']],
+            params_to_scrub=SETTINGS['scrub_fields'],
+        )  # priority: 50
     ]
+
+    # Add custom transforms
+    if len(SETTINGS['custom_transforms']) > 0:
+        _transforms.extend(SETTINGS['custom_transforms'])
+
+    # Sort the transforms by priority
+    _transforms = sorted(_transforms, key=lambda x: x.priority)
+
     _threads = queue.Queue()
     events.reset()
     filters.add_builtin_filters(SETTINGS)
diff --git a/rollbar/lib/transform.py b/rollbar/lib/transform.py
index 86b4e39..70b957b 100644
--- a/rollbar/lib/transform.py
+++ b/rollbar/lib/transform.py
@@ -1,5 +1,6 @@
 class Transform(object):
     depth_first = True
+    priority = 100
 
     def default(self, o, key=None):
         return o
diff --git a/rollbar/lib/transforms/scrub.py b/rollbar/lib/transforms/scrub.py
index ff0af81..e3b4264 100644
--- a/rollbar/lib/transforms/scrub.py
+++ b/rollbar/lib/transforms/scrub.py
@@ -6,6 +6,7 @@ from rollbar.lib.transform import Transform
 
 class ScrubTransform(Transform):
     suffix_matcher = None
+    priority = 40
     def __init__(self, suffixes=None, redact_char='*', randomize_len=True):
         super(ScrubTransform, self).__init__()
         if suffixes is not None and len(suffixes) > 0:
diff --git a/rollbar/lib/transforms/scrub_redact.py b/rollbar/lib/transforms/scrub_redact.py
index a81234b..09b9a76 100644
--- a/rollbar/lib/transforms/scrub_redact.py
+++ b/rollbar/lib/transforms/scrub_redact.py
@@ -9,6 +9,7 @@ REDACT_REF = RedactRef()
 
 
 class ScrubRedactTransform(ScrubTransform):
+    priority = 20
     def default(self, o, key=None):
         if o is REDACT_REF:
             return self.redact(o)
diff --git a/rollbar/lib/transforms/scruburl.py b/rollbar/lib/transforms/scruburl.py
index 0ab922c..a836614 100644
--- a/rollbar/lib/transforms/scruburl.py
+++ b/rollbar/lib/transforms/scruburl.py
@@ -9,6 +9,7 @@ _starts_with_auth_re = re.compile(r'^[a-zA-Z0-9-_]*(:[^@/]+)?@')
 
 
 class ScrubUrlTransform(ScrubTransform):
+    priority = 50
     def __init__(self,
                  suffixes=None,
                  scrub_username=False,
diff --git a/rollbar/lib/transforms/serializable.py b/rollbar/lib/transforms/serializable.py
index 49f95d2..f426f2d 100644
--- a/rollbar/lib/transforms/serializable.py
+++ b/rollbar/lib/transforms/serializable.py
@@ -9,6 +9,7 @@ from rollbar.lib.transform import Transform
 
 
 class SerializableTransform(Transform):
+    priority = 30
     def __init__(self, safe_repr=True, safelist_types=None):
         super(SerializableTransform, self).__init__()
         self.safe_repr = safe_repr
diff --git a/rollbar/lib/transforms/shortener.py b/rollbar/lib/transforms/shortener.py
index 383b163..1717673 100644
--- a/rollbar/lib/transforms/shortener.py
+++ b/rollbar/lib/transforms/shortener.py
@@ -108,6 +108,7 @@ def shorten_tuple(obj: tuple, max_len: int) -> tuple:
 
 class ShortenerTransform(Transform):
     depth_first = False
+    priority = 10
 
     def __init__(self, safe_repr=True, keys=None, **sizes):
         super(ShortenerTransform, self).__init__()


Test patch:
diff --git a/rollbar/test/test_transform.py b/rollbar/test/test_transform.py
new file mode 100644
index 0000000..3851997
--- /dev/null
+++ b/rollbar/test/test_transform.py
@@ -0,0 +1,93 @@
+import copy
+
+import rollbar
+from rollbar.test import BaseTest
+from rollbar.lib.transforms import Transform
+from rollbar.lib.transforms.shortener import ShortenerTransform
+from rollbar.lib.transforms.scrub_redact import ScrubRedactTransform
+from rollbar.lib.transforms.serializable import SerializableTransform
+from rollbar.lib.transforms.scrub import ScrubTransform
+from rollbar.lib.transforms.scruburl import ScrubUrlTransform
+
+_test_access_token = 'aaaabbbbccccddddeeeeffff00001111'
+_default_settings = copy.deepcopy(rollbar.SETTINGS)
+
+
+class TransformTest(BaseTest):
+    def setUp(self):
+        rollbar._initialized = False
+        rollbar.SETTINGS = copy.deepcopy(_default_settings)
+        rollbar.init(_test_access_token, locals={'enabled': True}, handler='blocking', timeout=12345)
+
+    def test_default_transforms(self):
+        transforms = {transform.__class__ for transform in rollbar._transforms}
+
+        self.assertEqual({
+            ShortenerTransform,
+            ScrubRedactTransform,
+            SerializableTransform,
+            ScrubTransform,
+            ScrubUrlTransform,
+        }, transforms)
+
+    def test_add_custom_transform(self):
+        class CustomTransform(Transform):
+            pass
+
+        rollbar._initialized = False
+        rollbar.SETTINGS = copy.deepcopy(_default_settings)
+        rollbar.init(_test_access_token,
+                     locals={'enabled': True},
+                     handler='blocking',
+                     timeout=12345,
+                     custom_transforms=[CustomTransform()])
+
+        transforms = {transform.__class__ for transform in rollbar._transforms}
+        transforms_ordered = [transform.__class__ for transform in rollbar._transforms]
+
+        self.assertEqual({
+            ShortenerTransform,
+            ScrubRedactTransform,
+            SerializableTransform,
+            ScrubTransform,
+            ScrubUrlTransform,
+            CustomTransform,
+        }, transforms)
+
+        # CustomTransform should be last because it has the default priority of 100
+        self.assertEqual([
+            ShortenerTransform,
+            ScrubRedactTransform,
+            SerializableTransform,
+            ScrubTransform,
+            ScrubUrlTransform,
+            CustomTransform,
+        ], transforms_ordered)
+
+    def test_add_custom_transform_first(self):
+        class CustomTransform(Transform):
+            priority = 1
+
+        class CustomTransformTwo(Transform):
+            priority = 35
+
+        rollbar._initialized = False
+        rollbar.SETTINGS = copy.deepcopy(_default_settings)
+        rollbar.init(_test_access_token,
+                     locals={'enabled': True},
+                     handler='blocking',
+                     timeout=12345,
+                     custom_transforms=[CustomTransform(), CustomTransformTwo()])
+
+        transforms = [transform.__class__ for transform in rollbar._transforms]
+
+        # Custom transforms should be first and fifth.
+        self.assertEqual([
+            CustomTransform,  # priority 1
+            ShortenerTransform,  # priority 10
+            ScrubRedactTransform,  # priority 20
+            SerializableTransform,  # priority 30
+            CustomTransformTwo,  # priority 35
+            ScrubTransform,  # priority 40
+            ScrubUrlTransform,  # priority 50
+        ], transforms)