ref(serializer): Make trimming faster by a few ms (getsentry#509)

* ref(serializer): Make trimming faster by a few ms

Refactor MetaNode and merge it into Serializer. 40% speedup for some
random microbenchmark, lol

* fix: Fix tests

Author: untitaker

sentry_sdk/serializer.py

@@ -1,4 +1,5 @@
 import contextlib
+import itertools
 
 from datetime import datetime
 
@@ -17,6 +18,7 @@
     from typing import Any
     from typing import Dict
     from typing import List
+    from typing import Tuple
     from typing import Optional
     from typing import Callable
     from typing import Union
@@ -50,81 +52,73 @@ def add_global_repr_processor(processor):
     global_repr_processors.append(processor)
 
 
-class MetaNode(object):
-    __slots__ = (
-        "_parent",
-        "_segment",
-        "_depth",
-        "_data",
-        "_is_databag",
-        "_should_repr_strings",
-    )
+class Memo(object):
+    def __init__(self):
+        # type: () -> None
+        self._inner = {}  # type: Dict[int, Any]
+
+    @contextlib.contextmanager
+    def memoize(self, obj):
+        # type: (Any) -> Generator[bool, None, None]
+        if id(obj) in self._inner:
+            yield True
+        else:
+            self._inner[id(obj)] = obj
+            yield False
+
+            self._inner.pop(id(obj), None)
+
+
+class Serializer(object):
+    __slots__ = ("memo", "_path", "_meta_stack", "_is_databag", "_should_repr_strings")
 
     def __init__(self):
         # type: () -> None
-        self._parent = None  # type: Optional[MetaNode]
-        self._segment = None  # type: Optional[Segment]
-        self._depth = 0  # type: int
-        self._data = None  # type: Optional[Dict[str, Any]]
+        self.memo = Memo()
+
+        self._path = []  # type: List[Segment]
+        self._meta_stack = []  # type: List[Dict[Segment, Any]]
         self._is_databag = None  # type: Optional[bool]
         self._should_repr_strings = None  # type: Optional[bool]
 
     def startswith_path(self, path):
-        # type: (List[Optional[str]]) -> bool
-        if len(path) > self._depth:
+        # type: (Tuple[Optional[Segment], ...]) -> bool
+        if len(path) > len(self._path):
             return False
 
-        return self.is_path(path + [None] * (self._depth - len(path)))
-
-    def is_path(self, path):
-        # type: (List[Optional[str]]) -> bool
-        if len(path) != self._depth:
-            return False
+        for i, segment in enumerate(path):
+            if segment is None:
+                continue
 
-        cur = self
-        for segment in reversed(path):
-            if segment is not None and segment != cur._segment:
+            if self._path[i] != segment:
                 return False
-            assert cur._parent is not None
-            cur = cur._parent
 
-        return cur._segment is None
-
-    def enter(self, segment):
-        # type: (Segment) -> MetaNode
-        rv = MetaNode()
-        rv._parent = self
-        rv._depth = self._depth + 1
-        rv._segment = segment
-        return rv
+        return True
 
-    def _create_annotations(self):
-        # type: () -> None
-        if self._data is not None:
-            return
+    def annotate(self, **meta):
+        # type: (**Any) -> None
+        while len(self._meta_stack) <= len(self._path):
+            try:
+                segment = self._path[len(self._meta_stack) - 1]
+                node = self._meta_stack[-1].setdefault(text_type(segment), {})
+            except IndexError:
+                node = {}
 
-        self._data = {}
-        if self._parent is not None:
-            self._parent._create_annotations()
-            self._parent._data[str(self._segment)] = self._data  # type: ignore
+            self._meta_stack.append(node)
 
-    def annotate(self, **meta):
-        # type: (Any) -> None
-        self._create_annotations()
-        assert self._data is not None
-        self._data.setdefault("", {}).update(meta)
+        self._meta_stack[-1].setdefault("", {}).update(meta)
 
     def should_repr_strings(self):
         # type: () -> bool
         if self._should_repr_strings is None:
             self._should_repr_strings = (
                 self.startswith_path(
-                    ["exception", "values", None, "stacktrace", "frames", None, "vars"]
+                    ("exception", "values", None, "stacktrace", "frames", None, "vars")
                 )
                 or self.startswith_path(
-                    ["threads", "values", None, "stacktrace", "frames", None, "vars"]
+                    ("threads", "values", None, "stacktrace", "frames", None, "vars")
                 )
-                or self.startswith_path(["stacktrace", "frames", None, "vars"])
+                or self.startswith_path(("stacktrace", "frames", None, "vars"))
             )
 
         return self._should_repr_strings
@@ -133,153 +127,120 @@ def is_databag(self):
         # type: () -> bool
         if self._is_databag is None:
             self._is_databag = (
-                self.startswith_path(["request", "data"])
-                or self.startswith_path(["breadcrumbs", None])
-                or self.startswith_path(["extra"])
-                or self.startswith_path(
-                    ["exception", "values", None, "stacktrace", "frames", None, "vars"]
-                )
-                or self.startswith_path(
-                    ["threads", "values", None, "stacktrace", "frames", None, "vars"]
-                )
-                or self.startswith_path(["stacktrace", "frames", None, "vars"])
+                self.should_repr_strings()
+                or self.startswith_path(("request", "data"))
+                or self.startswith_path(("breadcrumbs", None))
+                or self.startswith_path(("extra",))
             )
 
         return self._is_databag
 
-
-def _flatten_annotated(obj, meta_node):
-    # type: (Any, MetaNode) -> Any
-    if isinstance(obj, AnnotatedValue):
-        meta_node.annotate(**obj.metadata)
-        obj = obj.value
-    return obj
-
-
-class Memo(object):
-    def __init__(self):
-        # type: () -> None
-        self._inner = {}  # type: Dict[int, Any]
-
-    @contextlib.contextmanager
-    def memoize(self, obj):
-        # type: (Any) -> Generator[bool, None, None]
-        if id(obj) in self._inner:
-            yield True
-        else:
-            self._inner[id(obj)] = obj
-            yield False
-
-            self._inner.pop(id(obj), None)
-
-
-class Serializer(object):
-    def __init__(self):
-        # type: () -> None
-        self.memo = Memo()
-        self.meta_node = MetaNode()
-
-    @contextlib.contextmanager
-    def enter(self, segment):
-        # type: (Segment) -> Generator[None, None, None]
-        old_node = self.meta_node
-        self.meta_node = self.meta_node.enter(segment)
-
-        try:
-            yield
-        finally:
-            self.meta_node = old_node
-
     def serialize_event(self, obj):
         # type: (Any) -> Dict[str, Any]
         rv = self._serialize_node(obj)
-        if self.meta_node._data is not None:
-            rv["_meta"] = self.meta_node._data
+        if self._meta_stack:
+            rv["_meta"] = self._meta_stack[0]
         return rv
 
-    def _serialize_node(self, obj, max_depth=None, max_breadth=None):
-        # type: (Any, Optional[int], Optional[int]) -> Any
-        with capture_internal_exceptions():
-            with self.memo.memoize(obj) as result:
-                if result:
-                    return CYCLE_MARKER
+    def _serialize_node(self, obj, max_depth=None, max_breadth=None, segment=None):
+        # type: (Any, Optional[int], Optional[int], Optional[Segment]) -> Any
+        if segment is not None:
+            self._path.append(segment)
+            self._is_databag = self._is_databag or None
+            self._should_repr_strings = self._should_repr_strings or None
 
-                return self._serialize_node_impl(
-                    obj, max_depth=max_depth, max_breadth=max_breadth
-                )
+        try:
+            with capture_internal_exceptions():
+                with self.memo.memoize(obj) as result:
+                    if result:
+                        return CYCLE_MARKER
+
+                    return self._serialize_node_impl(
+                        obj, max_depth=max_depth, max_breadth=max_breadth
+                    )
 
-        if self.meta_node.is_databag():
-            return u"<failed to serialize, use init(debug=True) to see error logs>"
+            if self.is_databag():
+                return u"<failed to serialize, use init(debug=True) to see error logs>"
 
-        return None
+            return None
+        finally:
+            if segment is not None:
+                self._path.pop()
+                del self._meta_stack[len(self._path) + 1 :]
+                self._is_databag = self._is_databag and None
+                self._should_repr_strings = self._should_repr_strings and None
+
+    def _flatten_annotated(self, obj):
+        # type: (Any) -> Any
+        if isinstance(obj, AnnotatedValue):
+            self.annotate(**obj.metadata)
+            obj = obj.value
+        return obj
 
     def _serialize_node_impl(self, obj, max_depth, max_breadth):
         # type: (Any, Optional[int], Optional[int]) -> Any
-        if max_depth is None and max_breadth is None and self.meta_node.is_databag():
-            max_depth = self.meta_node._depth + MAX_DATABAG_DEPTH
-            max_breadth = self.meta_node._depth + MAX_DATABAG_BREADTH
+        cur_depth = len(self._path)
+        if max_depth is None and max_breadth is None and self.is_databag():
+            max_depth = cur_depth + MAX_DATABAG_DEPTH
+            max_breadth = cur_depth + MAX_DATABAG_BREADTH
 
         if max_depth is None:
             remaining_depth = None
         else:
-            remaining_depth = max_depth - self.meta_node._depth
+            remaining_depth = max_depth - cur_depth
 
-        obj = _flatten_annotated(obj, self.meta_node)
+        obj = self._flatten_annotated(obj)
 
         if remaining_depth is not None and remaining_depth <= 0:
-            self.meta_node.annotate(rem=[["!limit", "x"]])
-            if self.meta_node.is_databag():
-                return _flatten_annotated(strip_string(safe_repr(obj)), self.meta_node)
+            self.annotate(rem=[["!limit", "x"]])
+            if self.is_databag():
+                return self._flatten_annotated(strip_string(safe_repr(obj)))
             return None
 
-        if self.meta_node.is_databag():
+        if global_repr_processors and self.is_databag():
             hints = {"memo": self.memo, "remaining_depth": remaining_depth}
             for processor in global_repr_processors:
                 with capture_internal_exceptions():
                     result = processor(obj, hints)
                     if result is not NotImplemented:
-                        return _flatten_annotated(result, self.meta_node)
+                        return self._flatten_annotated(result)
 
         if isinstance(obj, Mapping):
-            # Create temporary list here to avoid calling too much code that
+            # Create temporary copy here to avoid calling too much code that
             # might mutate our dictionary while we're still iterating over it.
-            items = []
-            for i, (k, v) in enumerate(iteritems(obj)):
-                if max_breadth is not None and i >= max_breadth:
-                    self.meta_node.annotate(len=max_breadth)
-                    break
-
-                items.append((k, v))
-
-            rv_dict = {}  # type: Dict[Any, Any]
-            for k, v in items:
-                k = text_type(k)
-
-                with self.enter(k):
-                    v = self._serialize_node(
-                        v, max_depth=max_depth, max_breadth=max_breadth
-                    )
-                    if v is not None:
-                        rv_dict[k] = v
+            if max_breadth is not None and len(obj) >= max_breadth:
+                rv_dict = dict(itertools.islice(iteritems(obj), None, max_breadth))
+                self.annotate(len=len(obj))
+            else:
+                rv_dict = dict(iteritems(obj))
+
+            for k in list(rv_dict):
+                str_k = text_type(k)
+                v = self._serialize_node(
+                    rv_dict.pop(k),
+                    max_depth=max_depth,
+                    max_breadth=max_breadth,
+                    segment=str_k,
+                )
+                if v is not None:
+                    rv_dict[str_k] = v
 
             return rv_dict
-        elif isinstance(obj, Sequence) and not isinstance(obj, string_types):
-            rv_list = []  # type: List[Any]
-            for i, v in enumerate(obj):
-                if max_breadth is not None and i >= max_breadth:
-                    self.meta_node.annotate(len=max_breadth)
-                    break
-
-                with self.enter(i):
-                    rv_list.append(
-                        self._serialize_node(
-                            v, max_depth=max_depth, max_breadth=max_breadth
-                        )
-                    )
+        elif not isinstance(obj, string_types) and isinstance(obj, Sequence):
+            if max_breadth is not None and len(obj) >= max_breadth:
+                rv_list = list(obj)[:max_breadth]
+                self.annotate(len=len(obj))
+            else:
+                rv_list = list(obj)
+
+            for i in range(len(rv_list)):
+                rv_list[i] = self._serialize_node(
+                    rv_list[i], max_depth=max_depth, max_breadth=max_breadth, segment=i
+                )
 
             return rv_list
 
-        if self.meta_node.should_repr_strings():
+        if self.should_repr_strings():
             obj = safe_repr(obj)
         else:
             if obj is None or isinstance(obj, (bool, number_types)):
@@ -294,4 +255,4 @@ def _serialize_node_impl(self, obj, max_depth, max_breadth):
             if not isinstance(obj, string_types):
                 obj = safe_repr(obj)
 
-        return _flatten_annotated(strip_string(obj), self.meta_node)
+        return self._flatten_annotated(strip_string(obj))