Merge pull request #461 from rollbar/changed/shortner-to-breadth-first-traverse

danielmorell · web-flow · commit 9d0857764caf · 2024-10-29T09:28:06.000-05:00
Changed the `ShortenerTransform` to use breadth first traversal
diff --git a/rollbar/lib/transform.py b/rollbar/lib/transform.py
@@ -1,4 +1,6 @@
 class Transform(object):
+    depth_first = True
+
     def default(self, o, key=None):
         return o
 
diff --git a/rollbar/lib/transforms/__init__.py b/rollbar/lib/transforms/__init__.py
@@ -89,7 +89,7 @@ def default_handler(o, key=None):
         "allowed_circular_reference_types": _ALLOWED_CIRCULAR_REFERENCE_TYPES,
     }
 
-    return traverse.traverse(obj, key=key, **handlers)
+    return traverse.traverse(obj, key=key, depth_first=transform.depth_first, **handlers)
 
 
 __all__ = ["transform", "Transform"]
diff --git a/rollbar/lib/transforms/shortener.py b/rollbar/lib/transforms/shortener.py
@@ -4,9 +4,10 @@
 import reprlib
 
 from collections.abc import Mapping
+from typing import Union, Tuple
 
 from rollbar.lib import (
-    integer_types, key_in, key_depth, number_types, sequence_types,
+    integer_types, key_in, key_depth, sequence_types,
     string_types)
 from rollbar.lib.transform import Transform
 
@@ -25,7 +26,89 @@
 }
 
 
+def _max_left_right(max_len: int, seperator_len: int) -> Tuple[int, int]:
+    left = max(0, (max_len-seperator_len)//2)
+    right = max(0, max_len-seperator_len-left)
+    return left, right
+
+
+def shorten_array(obj: array, max_len: int) -> array:
+    if len(obj) <= max_len:
+        return obj
+
+    return obj[:max_len]
+
+
+def shorten_bytes(obj: bytes, max_len: int) -> bytes:
+    if len(obj) <= max_len:
+        return obj
+
+    return obj[:max_len]
+
+
+def shorten_deque(obj: collections.deque, max_len: int) -> collections.deque:
+    if len(obj) <= max_len:
+        return obj
+
+    return collections.deque(itertools.islice(obj, max_len))
+
+
+def shorten_frozenset(obj: frozenset, max_len: int) -> frozenset:
+    if len(obj) <= max_len:
+        return obj
+
+    return frozenset([elem for i, elem in enumerate(obj) if i < max_len] + ['...'])
+
+
+def shorten_int(obj: int, max_len: int) -> Union[int, str]:
+    s = repr(obj)
+    if len(s) <= max_len:
+        return obj
+
+    left, right = _max_left_right(max_len, 3)
+    return s[:left] + '...' + s[len(s)-right:]
+
+
+def shorten_list(obj: list, max_len: int) -> list:
+    if len(obj) <= max_len:
+        return obj
+
+    return obj[:max_len] + ['...']
+
+
+def shorten_mapping(obj: Union[dict, Mapping], max_keys: int) -> dict:
+    if len(obj) <= max_keys:
+        return obj
+
+    return {k: obj[k] for k in itertools.islice(obj.keys(), max_keys)}
+
+
+def shorten_set(obj: set, max_len: int) -> set:
+    if len(obj) <= max_len:
+        return obj
+
+    return set([elem for i, elem in enumerate(obj) if i < max_len] + ['...'])
+
+
+def shorten_string(obj: str, max_len: int) -> str:
+    if len(obj) <= max_len:
+        return obj
+
+    left, right = _max_left_right(max_len, 3)
+    return obj[:left] + '...' + obj[len(obj)-right:]
+
+
+def shorten_tuple(obj: tuple, max_len: int) -> tuple:
+    if len(obj) <= max_len:
+        return obj
+
+    return obj[:max_len] + ('...',)
+
+
+
 class ShortenerTransform(Transform):
+    depth_first = False
+
     def __init__(self, safe_repr=True, keys=None, **sizes):
         super(ShortenerTransform, self).__init__()
         self.safe_repr = safe_repr
@@ -47,26 +130,33 @@ def _get_max_size(self, obj):
 
         return self._repr.maxother
 
-    def _shorten_sequence(self, obj, max_keys):
-        _len = len(obj)
-        if _len <= max_keys:
-            return obj
-
-        return self._repr.repr(obj)
-
-    def _shorten_mapping(self, obj, max_keys):
-        _len = len(obj)
-        if _len <= max_keys:
-            return obj
-
-        return {k: obj[k] for k in itertools.islice(obj.keys(), max_keys)}
+    def _shorten(self, val):
+        max_size = self._get_max_size(val)
 
-    def _shorten_basic(self, obj, max_len):
-        val = str(obj)
-        if len(val) <= max_len:
-            return obj
+        if isinstance(val, array):
+            return shorten_array(val, max_size)
+        if isinstance(val, bytes):
+            return shorten_bytes(val, max_size)
+        if isinstance(val, collections.deque):
+            return shorten_deque(val, max_size)
+        if isinstance(val, (dict, Mapping)):
+            return shorten_mapping(val, max_size)
+        if isinstance(val, float):
+            return val
+        if isinstance(val, frozenset):
+            return shorten_frozenset(val, max_size)
+        if isinstance(val, int):
+            return shorten_int(val, max_size)
+        if isinstance(val, list):
+            return shorten_list(val, max_size)
+        if isinstance(val, set):
+            return shorten_set(val, max_size)
+        if isinstance(val, str):
+            return shorten_string(val, max_size)
+        if isinstance(val, tuple):
+            return shorten_tuple(val, max_size)
 
-        return self._repr.repr(obj)
+        return self._shorten_other(val)
 
     def _shorten_other(self, obj):
         if obj is None:
@@ -77,19 +167,6 @@ def _shorten_other(self, obj):
 
         return self._repr.repr(obj)
 
-    def _shorten(self, val):
-        max_size = self._get_max_size(val)
-
-        if isinstance(val, dict):
-            return self._shorten_mapping(val, max_size)
-        if isinstance(val, (string_types, sequence_types)):
-            return self._shorten_sequence(val, max_size)
-
-        if isinstance(val, number_types):
-            return self._shorten_basic(val, self._repr.maxlong)
-
-        return self._shorten_other(val)
-
     def _should_shorten(self, val, key):
         if not key:
             return False
@@ -100,18 +177,18 @@ def _should_drop(self, val, key) -> bool:
         if not key:
             return False
 
-        maxdepth = key_depth(key, self.keys)
-        if maxdepth == 0:
+        max_depth = key_depth(key, self.keys)
+        if max_depth == 0:
             return False
 
-        return (maxdepth + self._repr.maxlevel) <= len(key)
+        return (max_depth + self._repr.maxlevel) <= len(key)
 
     def default(self, o, key=None):
         if self._should_drop(o, key):
-            if isinstance(o, dict):
-                return '{...}'
+            if isinstance(o, (dict, Mapping)):
+                return {'...': '...'}
             if isinstance(o, sequence_types):
-                return '[...]'
+                return ['...']
 
         if self._should_shorten(o, key):
             return self._shorten(o)
diff --git a/rollbar/lib/traverse.py b/rollbar/lib/traverse.py
@@ -81,6 +81,7 @@ def traverse(
     circular_reference_handler=_default_handlers[CIRCULAR],
     allowed_circular_reference_types=None,
     memo=None,
+    depth_first=True,
     **custom_handlers
 ):
     memo = memo or {}
@@ -108,42 +109,58 @@ def traverse(
         "circular_reference_handler": circular_reference_handler,
         "allowed_circular_reference_types": allowed_circular_reference_types,
         "memo": memo,
+        "depth_first": depth_first,
     }
     kw.update(custom_handlers)
 
     try:
         if obj_type is STRING:
             return string_handler(obj, key=key)
         elif obj_type is TUPLE:
-            return tuple_handler(
-                tuple(
-                    traverse(elem, key=key + (i,), **kw) for i, elem in enumerate(obj)
-                ),
-                key=key,
-            )
+            if depth_first:
+                return tuple_handler(
+                    tuple(
+                        traverse(elem, key=key + (i,), **kw) for i, elem in enumerate(obj)
+                    ),
+                    key=key,
+                )
+            # Breadth first
+            return tuple(traverse(elem, key=key + (i,), **kw) for i, elem in enumerate(tuple_handler(obj, key=key)))
         elif obj_type is NAMEDTUPLE:
-            return namedtuple_handler(
-                obj._make(
-                    traverse(v, key=key + (k,), **kw)
-                    for k, v in obj._asdict().items()
-                ),
-                key=key,
-            )
+            if depth_first:
+                return namedtuple_handler(
+                    obj._make(
+                        traverse(v, key=key + (k,), **kw)
+                        for k, v in obj._asdict().items()
+                    ),
+                    key=key,
+                )
+            # Breadth first
+            return obj._make(traverse(v, key=key + (k,), **kw) for k, v in namedtuple_handler(obj, key=key)._asdict().items())
         elif obj_type is LIST:
-            return list_handler(
-                [traverse(elem, key=key + (i,), **kw) for i, elem in enumerate(obj)],
-                key=key,
-            )
+            if depth_first:
+                return list_handler(
+                    [traverse(elem, key=key + (i,), **kw) for i, elem in enumerate(obj)],
+                    key=key,
+                )
+            # Breadth first
+            return [traverse(elem, key=key + (i,), **kw) for i, elem in enumerate(list_handler(obj, key=key))]
         elif obj_type is SET:
-            return set_handler(
-                {traverse(elem, key=key + (i,), **kw) for i, elem in enumerate(obj)},
-                key=key,
-            )
+            if depth_first:
+                return set_handler(
+                    {traverse(elem, key=key + (i,), **kw) for i, elem in enumerate(obj)},
+                    key=key,
+                )
+            # Breadth first
+            return {traverse(elem, key=key + (i,), **kw) for i, elem in enumerate(set_handler(obj, key=key))}
         elif obj_type is MAPPING:
-            return mapping_handler(
-                {k: traverse(v, key=key + (k,), **kw) for k, v in obj.items()},
-                key=key,
-            )
+            if depth_first:
+                return mapping_handler(
+                    {k: traverse(v, key=key + (k,), **kw) for k, v in obj.items()},
+                    key=key,
+                )
+            # Breadth first
+            return {k: traverse(v, key=key + (k,), **kw) for k, v in mapping_handler(obj, key=key).items()}
         elif obj_type is PATH:
             return path_handler(obj, key=key)
         elif obj_type is DEFAULT:
diff --git a/rollbar/test/test_rollbar.py b/rollbar/test/test_rollbar.py
@@ -1555,7 +1555,7 @@ def _raise(large):
 
         self.assertEqual(1, len(payload['data']['body']['trace']['frames'][-1]['argspec']))
         self.assertEqual('large', payload['data']['body']['trace']['frames'][-1]['argspec'][0])
-        self.assertEqual("'###############################################...################################################'",
+        self.assertEqual("################################################...#################################################",
                          payload['data']['body']['trace']['frames'][-1]['locals']['large'])
 
     @mock.patch('rollbar.send_payload')
@@ -1586,12 +1586,12 @@ def _raise(large):
 
         self.assertEqual('large', payload['data']['body']['trace']['frames'][-1]['argspec'][0])
         self.assertTrue(
-            ("['hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', ...]" ==
+            (['hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', '...'] ==
                 payload['data']['body']['trace']['frames'][-1]['argspec'][0])
 
             or
 
-            ("['hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', ...]" ==
+            (['hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', 'hi', '...'] ==
                     payload['data']['body']['trace']['frames'][0]['locals']['xlarge']))
 
 
diff --git a/rollbar/test/test_shortener_transform.py b/rollbar/test/test_shortener_transform.py
diff --git a/rollbar/test/test_traverse.py b/rollbar/test/test_traverse.py

Original file line number	Diff line number	Diff line change
`@@ -89,7 +89,7 @@ def default_handler(o, key=None):`
`89`	`89`	`"allowed_circular_reference_types": _ALLOWED_CIRCULAR_REFERENCE_TYPES,`
`90`	`90`	`}`
`91`	`91`
`92`		`- return traverse.traverse(obj, key=key, **handlers)`
	`92`	`+ return traverse.traverse(obj, key=key, depth_first=transform.depth_first, **handlers)`
`93`	`93`
`94`	`94`
`95`	`95`	`__all__ = ["transform", "Transform"]`