feat: change config name, add some tests, add limiter

Author: ablaszkiewicz

posthog/__init__.py

@@ -10,7 +10,7 @@
     tag as inner_tag,
     set_context_session as inner_set_context_session,
     identify_context as inner_identify_context,
-    set_code_variables_enabled_context as inner_set_code_variables_enabled_context,
+    set_capture_exception_code_variables_context as inner_set_capture_exception_code_variables_context,
     set_code_variables_mask_patterns_context as inner_set_code_variables_mask_patterns_context,
     set_code_variables_ignore_patterns_context as inner_set_code_variables_ignore_patterns_context,
 )
@@ -109,11 +109,11 @@ def identify_context(distinct_id: str):
     return inner_identify_context(distinct_id)
 
 
-def set_code_variables_enabled_context(enabled: bool):
+def set_capture_exception_code_variables_context(enabled: bool):
     """
     Set whether code variables are captured for the current context.
     """
-    return inner_set_code_variables_enabled_context(enabled)
+    return inner_set_capture_exception_code_variables_context(enabled)
 
 
 def set_code_variables_mask_patterns_context(mask_patterns: list):

posthog/client.py

@@ -42,7 +42,7 @@
     _get_current_context,
     get_context_distinct_id,
     get_context_session_id,
-    get_code_variables_enabled_context,
+    get_capture_exception_code_variables_context,
     get_code_variables_mask_patterns_context,
     get_code_variables_ignore_patterns_context,
     new_context,
@@ -185,7 +185,7 @@ def __init__(
         before_send=None,
         flag_fallback_cache_url=None,
         enable_local_evaluation=True,
-        code_variables_enabled=False,
+        capture_exception_code_variables=False,
         code_variables_mask_patterns=None,
         code_variables_ignore_patterns=None,
     ):
@@ -243,8 +243,7 @@ def __init__(
         self.privacy_mode = privacy_mode
         self.enable_local_evaluation = enable_local_evaluation
 
-        # Code variables configuration
-        self.code_variables_enabled = code_variables_enabled
+        self.capture_exception_code_variables = capture_exception_code_variables
         self.code_variables_mask_patterns = (
             code_variables_mask_patterns
             if code_variables_mask_patterns is not None
@@ -1002,14 +1001,14 @@ def capture_exception(
                 **properties,
             }
 
-            context_enabled = get_code_variables_enabled_context()
+            context_enabled = get_capture_exception_code_variables_context()
             context_mask = get_code_variables_mask_patterns_context()
             context_ignore = get_code_variables_ignore_patterns_context()
 
             enabled = (
                 context_enabled
                 if context_enabled is not None
-                else self.code_variables_enabled
+                else self.capture_exception_code_variables
             )
             mask_patterns = (
                 context_mask

posthog/contexts.py

@@ -22,7 +22,7 @@ def __init__(
         self.session_id: Optional[str] = None
         self.distinct_id: Optional[str] = None
         self.tags: Dict[str, Any] = {}
-        self.code_variables_enabled: Optional[bool] = None
+        self.capture_exception_code_variables: Optional[bool] = None
         self.code_variables_mask_patterns: Optional[list] = None
         self.code_variables_ignore_patterns: Optional[list] = None
 
@@ -35,8 +35,8 @@ def set_distinct_id(self, distinct_id: str):
     def add_tag(self, key: str, value: Any):
         self.tags[key] = value
 
-    def set_code_variables_enabled(self, enabled: bool):
-        self.code_variables_enabled = enabled
+    def set_capture_exception_code_variables(self, enabled: bool):
+        self.capture_exception_code_variables = enabled
 
     def set_code_variables_mask_patterns(self, mask_patterns: list):
         self.code_variables_mask_patterns = mask_patterns
@@ -71,11 +71,11 @@ def collect_tags(self) -> Dict[str, Any]:
             tags.update(new_tags)
         return tags
 
-    def get_code_variables_enabled(self) -> Optional[bool]:
-        if self.code_variables_enabled is not None:
-            return self.code_variables_enabled
+    def get_capture_exception_code_variables(self) -> Optional[bool]:
+        if self.capture_exception_code_variables is not None:
+            return self.capture_exception_code_variables
         if self.parent is not None and not self.fresh:
-            return self.parent.get_code_variables_enabled()
+            return self.parent.get_capture_exception_code_variables()
         return None
 
     def get_code_variables_mask_patterns(self) -> Optional[list]:
@@ -276,13 +276,13 @@ def get_context_distinct_id() -> Optional[str]:
     return None
 
 
-def set_code_variables_enabled_context(enabled: bool) -> None:
+def set_capture_exception_code_variables_context(enabled: bool) -> None:
     """
     Set whether code variables are captured for the current context.
     """
     current_context = _get_current_context()
     if current_context:
-        current_context.set_code_variables_enabled(enabled)
+        current_context.set_capture_exception_code_variables(enabled)
 
 
 def set_code_variables_mask_patterns_context(mask_patterns: list) -> None:
@@ -303,10 +303,10 @@ def set_code_variables_ignore_patterns_context(ignore_patterns: list) -> None:
         current_context.set_code_variables_ignore_patterns(ignore_patterns)
 
 
-def get_code_variables_enabled_context() -> Optional[bool]:
+def get_capture_exception_code_variables_context() -> Optional[bool]:
     current_context = _get_current_context()
     if current_context:
-        return current_context.get_code_variables_enabled()
+        return current_context.get_capture_exception_code_variables()
     return None
 
 

posthog/exception_utils.py

@@ -60,6 +60,24 @@
 
 CODE_VARIABLES_REDACTED_VALUE = "$$_posthog_redacted_based_on_masking_rules_$$"
 
+DEFAULT_TOTAL_VARIABLES_SIZE_LIMIT = 20 * 1024
+
+
+class VariableSizeLimiter:
+    def __init__(self, max_size=DEFAULT_TOTAL_VARIABLES_SIZE_LIMIT):
+        self.max_size = max_size
+        self.current_size = 0
+
+    def can_add(self, size):
+        return self.current_size + size <= self.max_size
+
+    def add(self, size):
+        self.current_size += size
+
+    def get_remaining_space(self):
+        return self.max_size - self.current_size
+
+
 LogLevelStr = Literal["fatal", "critical", "error", "warning", "info", "debug"]
 
 Event = TypedDict(
@@ -923,37 +941,51 @@ def _pattern_matches(name, patterns):
     return False
 
 
-def _serialize_variable_value(value, max_length=1024):
+def _serialize_variable_value(value, limiter, max_length=1024):
     try:
         if value is None:
-            return "None"
+            result = "None"
         elif isinstance(value, bool):
-            return str(value)
+            result = str(value)
         elif isinstance(value, (int, float)):
             result = str(value)
         elif isinstance(value, str):
             result = repr(value)
         else:
             result = json.dumps(value)
 
-        # Truncate if too long
         if len(result) > max_length:
             result = result[: max_length - 3] + "..."
 
+        result_size = len(result)
+        if not limiter.can_add(result_size):
+            return None
+        limiter.add(result_size)
+
         return result
     except Exception:
         try:
-            return f"<{type(value).__name__}>"
+            fallback = f"<{type(value).__name__}>"
+            fallback_size = len(fallback)
+            if not limiter.can_add(fallback_size):
+                return None
+            limiter.add(fallback_size)
+            return fallback
         except Exception:
-            return "<unserializable object>"
+            fallback = "<unserializable object>"
+            fallback_size = len(fallback)
+            if not limiter.can_add(fallback_size):
+                return None
+            limiter.add(fallback_size)
+            return fallback
 
 
 def _is_simple_type(value):
     return isinstance(value, (type(None), bool, int, float, str))
 
 
 def serialize_code_variables(
-    frame, mask_patterns=None, ignore_patterns=None, max_vars=20, max_length=1024
+    frame, limiter, mask_patterns=None, ignore_patterns=None, max_length=1024
 ):
     if mask_patterns is None:
         mask_patterns = []
@@ -981,23 +1013,25 @@ def serialize_code_variables(
             complex_vars[name] = value
 
     result = {}
-    count = 0
 
     all_vars = {**simple_vars, **complex_vars}
     ordered_names = list(sorted(simple_vars.keys())) + list(sorted(complex_vars.keys()))
 
     for name in ordered_names:
-        if count >= max_vars:
-            break
-
         value = all_vars[name]
 
         if _pattern_matches(name, compiled_mask):
-            result[name] = CODE_VARIABLES_REDACTED_VALUE
+            redacted_value = CODE_VARIABLES_REDACTED_VALUE
+            redacted_size = len(redacted_value)
+            if not limiter.can_add(redacted_size):
+                break
+            limiter.add(redacted_size)
+            result[name] = redacted_value
         else:
-            result[name] = _serialize_variable_value(value, max_length)
-
-        count += 1
+            serialized = _serialize_variable_value(value, limiter, max_length)
+            if serialized is None:
+                break
+            result[name] = serialized
 
     return result
 
@@ -1015,6 +1049,8 @@ def try_attach_code_variables_to_frames(
     if not tb_frames:
         return
 
+    limiter = VariableSizeLimiter()
+
     for exception in all_exceptions:
         stacktrace = exception.get("stacktrace")
         if not stacktrace or "frames" not in stacktrace:
@@ -1028,9 +1064,9 @@ def try_attach_code_variables_to_frames(
 
             variables = serialize_code_variables(
                 tb_item.tb_frame,
+                limiter,
                 mask_patterns=mask_patterns,
                 ignore_patterns=ignore_patterns,
-                max_vars=20,
                 max_length=1024,
             )