Merge pull request #448 from rollbar/pawel/traverse_dict

traverse object for recursive shortening

rollbar/lib/transforms/shortener.py

@@ -47,6 +47,8 @@ def _get_max_size(self, obj):
 
         return self._repr.maxother
 
+    def _get_max_level(self):
+        return getattr(self._repr, 'maxlevel')
     def _shorten_sequence(self, obj, max_keys):
         _len = len(obj)
         if _len <= max_keys:
@@ -77,14 +79,44 @@ def _shorten_other(self, obj):
 
         return self._repr.repr(obj)
 
+    def traverse_obj(self, obj, level=1):
+        def seq_iter(o):
+            return o if isinstance(o, dict) else range(len(o))
+
+        for k in seq_iter(obj):
+            max_size = self._get_max_size(obj[k])
+            if isinstance(obj[k], dict):
+                obj[k] = self._shorten_mapping(obj[k], max_size)
+                if level == self._get_max_level():
+                    del obj[k]
+                    return
+                self.traverse_obj(obj[k], level + 1)
+            elif isinstance(obj[k], sequence_types):
+                obj[k] = self._shorten_sequence(obj[k], max_size)
+                if level == self._get_max_level():
+                    del obj[k]
+                    return
+                self.traverse_obj(obj[k], level + 1)
+            else:
+                obj[k] = self._shorten(obj[k])
+        return 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)):
+            val = self._shorten_mapping(val, max_size)
+            return self.traverse_obj(val)
+
+        if isinstance(val, string_types):
             return self._shorten_sequence(val, max_size)
 
+        if isinstance(val, sequence_types):
+            val = self._shorten_sequence(val, max_size)
+            if isinstance(val, string_types):
+                return val
+            return self.traverse_obj(val)
+
         if isinstance(val, number_types):
             return self._shorten_basic(val, self._repr.maxlong)
 

rollbar/test/test_shortener_transform.py

@@ -37,22 +37,27 @@ def setUp(self):
             'frozenset': frozenset([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]),
             'array': array('l', [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]),
             'deque': deque([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], 15),
-            'other': TestClassWithAVeryVeryVeryVeryVeryVeryVeryLongName()
+            'other': TestClassWithAVeryVeryVeryVeryVeryVeryVeryLongName(),
+            'list_max_level': [1, [2, [3, [4, ["good_5", ["bad_6", ["bad_7"]]]]]]],
+            'dict_max_level': {1: 1, 2: {3: {4: {"level4": "good", "level5": {"toplevel": "ok", 6: {7: {}}}}}}},
+            'list_multi_level':  [1, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]]
         }
 
     def _assert_shortened(self, key, expected):
         shortener = ShortenerTransform(keys=[(key,)], **DEFAULT_LOCALS_SIZES)
         result = transforms.transform(self.data, shortener)
 
         if key == 'dict':
-            self.assertEqual(expected, len(result))
+            self.assertEqual(expected, len(result[key]))
+        elif key in ('list_max_level', 'dict_max_level', 'list_multi_level'):
+            self.assertEqual(expected,  result[key])
         else:
             # the repr output can vary between Python versions
             stripped_result_key = result[key].strip("'\"u")
 
         if key == 'other':
             self.assertIn(expected, stripped_result_key)
-        elif key != 'dict':
+        elif key not in ('dict', 'list_max_level', 'dict_max_level', 'list_multi_level'):
             self.assertEqual(expected, stripped_result_key)
 
         # make sure nothing else was shortened
@@ -82,6 +87,18 @@ def test_shorten_list(self):
         expected = '[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, ...]'
         self._assert_shortened('list', expected)
 
+    def test_shorten_list_max_level(self):
+        expected = [1, [2, [3, [4, ['good_5']]]]]
+        self._assert_shortened('list_max_level', expected)
+
+    def test_shorten_list_multi_level(self):
+        expected = [1, '[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, ...]']
+        self._assert_shortened('list_multi_level', expected)
+
+    def test_shorten_dict_max_level(self):
+        expected = {1: 1, 2: {3: {4: {'level4': 'good', 'level5': {'toplevel': 'ok'}}}}}
+        self._assert_shortened('dict_max_level', expected)
+
     def test_shorten_tuple(self):
         expected = '(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, ...)'
         self._assert_shortened('tuple', expected)