feat: add code variables capture

Author: ablaszkiewicz

posthog/__init__.py

@@ -10,6 +10,9 @@
     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_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,
 )
 from posthog.feature_flags import InconclusiveMatchError, RequiresServerEvaluation
 from posthog.types import FeatureFlag, FlagsAndPayloads, FeatureFlagResult
@@ -103,6 +106,27 @@ def identify_context(distinct_id: str):
     return inner_identify_context(distinct_id)
 
 
+def set_code_variables_enabled_context(enabled: bool):
+    """
+    Set whether code variables are captured for the current context.
+    """
+    return inner_set_code_variables_enabled_context(enabled)
+
+
+def set_code_variables_mask_patterns_context(mask_patterns: list):
+    """
+    Variable names matching these patterns will be masked with *** when capturing code variables.
+    """
+    return inner_set_code_variables_mask_patterns_context(mask_patterns)
+
+
+def set_code_variables_ignore_patterns_context(ignore_patterns: list):
+    """
+    Variable names matching these patterns will be ignored completely when capturing code variables.
+    """
+    return inner_set_code_variables_ignore_patterns_context(ignore_patterns)
+
+
 def tag(name: str, value: Any):
     """
     Add a tag to the current context.

posthog/client.py

@@ -19,6 +19,9 @@
     handle_in_app,
     exception_is_already_captured,
     mark_exception_as_captured,
+    extract_code_variables_from_exception,
+    DEFAULT_CODE_VARIABLES_MASK_PATTERNS,
+    DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS,
 )
 from posthog.feature_flags import (
     InconclusiveMatchError,
@@ -39,6 +42,9 @@
     _get_current_context,
     get_context_distinct_id,
     get_context_session_id,
+    get_code_variables_enabled_context,
+    get_code_variables_mask_patterns_context,
+    get_code_variables_ignore_patterns_context,
     new_context,
 )
 from posthog.types import (
@@ -179,6 +185,9 @@ def __init__(
         before_send=None,
         flag_fallback_cache_url=None,
         enable_local_evaluation=True,
+        code_variables_enabled=False,
+        code_variables_mask_patterns=None,
+        code_variables_ignore_patterns=None,
     ):
         """
         Initialize a new PostHog client instance.
@@ -233,6 +242,17 @@ def __init__(
         self.exception_capture = None
         self.privacy_mode = privacy_mode
         self.enable_local_evaluation = enable_local_evaluation
+        
+        # Code variables configuration
+        self.code_variables_enabled = code_variables_enabled
+        self.code_variables_mask_patterns = (
+            code_variables_mask_patterns if code_variables_mask_patterns is not None 
+            else DEFAULT_CODE_VARIABLES_MASK_PATTERNS
+        )
+        self.code_variables_ignore_patterns = (
+            code_variables_ignore_patterns if code_variables_ignore_patterns is not None 
+            else DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS
+        )
 
         if project_root is None:
             try:
@@ -980,6 +1000,23 @@ def capture_exception(
                 **properties,
             }
 
+            context_enabled = get_code_variables_enabled_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
+            mask_patterns = context_mask if context_mask is not None else self.code_variables_mask_patterns
+            ignore_patterns = context_ignore if context_ignore is not None else self.code_variables_ignore_patterns
+            
+            if enabled:
+                code_variables = extract_code_variables_from_exception(
+                    exc_info,
+                    mask_patterns=mask_patterns,
+                    ignore_patterns=ignore_patterns
+                )
+                if code_variables:
+                    properties["$exception_code_variables"] = code_variables
+
             if self.log_captured_exceptions:
                 self.log.exception(exception, extra=kwargs)
 

posthog/contexts.py

@@ -22,6 +22,9 @@ 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.code_variables_mask_patterns: Optional[list] = None
+        self.code_variables_ignore_patterns: Optional[list] = None
 
     def set_session_id(self, session_id: str):
         self.session_id = session_id
@@ -32,6 +35,15 @@ 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_code_variables_mask_patterns(self, mask_patterns: list):
+        self.code_variables_mask_patterns = mask_patterns
+
+    def set_code_variables_ignore_patterns(self, ignore_patterns: list):
+        self.code_variables_ignore_patterns = ignore_patterns
+
     def get_parent(self):
         return self.parent
 
@@ -59,6 +71,27 @@ 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
+        if self.parent is not None and not self.fresh:
+            return self.parent.get_code_variables_enabled()
+        return None
+
+    def get_code_variables_mask_patterns(self) -> Optional[list]:
+        if self.code_variables_mask_patterns is not None:
+            return self.code_variables_mask_patterns
+        if self.parent is not None and not self.fresh:
+            return self.parent.get_code_variables_mask_patterns()
+        return None
+
+    def get_code_variables_ignore_patterns(self) -> Optional[list]:
+        if self.code_variables_ignore_patterns is not None:
+            return self.code_variables_ignore_patterns
+        if self.parent is not None and not self.fresh:
+            return self.parent.get_code_variables_ignore_patterns()
+        return None
+
 
 _context_stack: contextvars.ContextVar[Optional[ContextScope]] = contextvars.ContextVar(
     "posthog_context_stack", default=None
@@ -243,6 +276,54 @@ def get_context_distinct_id() -> Optional[str]:
     return None
 
 
+def set_code_variables_enabled_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)
+
+
+def set_code_variables_mask_patterns_context(mask_patterns: list) -> None:
+    """
+    Variable names matching these patterns will be masked with *** when capturing code variables.
+    """
+    current_context = _get_current_context()
+    if current_context:
+        current_context.set_code_variables_mask_patterns(mask_patterns)
+
+
+def set_code_variables_ignore_patterns_context(ignore_patterns: list) -> None:
+    """
+    Variable names matching these patterns will be ignored completely when capturing code variables.
+    """
+    current_context = _get_current_context()
+    if current_context:
+        current_context.set_code_variables_ignore_patterns(ignore_patterns)
+
+
+def get_code_variables_enabled_context() -> Optional[bool]:
+    current_context = _get_current_context()
+    if current_context:
+        return current_context.get_code_variables_enabled()
+    return None
+
+
+def get_code_variables_mask_patterns_context() -> Optional[list]:
+    current_context = _get_current_context()
+    if current_context:
+        return current_context.get_code_variables_mask_patterns()
+    return None
+
+
+def get_code_variables_ignore_patterns_context() -> Optional[list]:
+    current_context = _get_current_context()
+    if current_context:
+        return current_context.get_code_variables_ignore_patterns()
+    return None
+
+
 F = TypeVar("F", bound=Callable[..., Any])
 
 

posthog/exception_utils.py

@@ -26,6 +26,7 @@
     Union,
     cast,
     TYPE_CHECKING,
+    Pattern,
 )
 
 from posthog.args import ExcInfo, ExceptionArg  # noqa: F401
@@ -40,6 +41,22 @@
 
 DEFAULT_MAX_VALUE_LENGTH = 1024
 
+DEFAULT_CODE_VARIABLES_MASK_PATTERNS = [
+    r"(?i).*password.*",
+    r"(?i).*secret.*",
+    r"(?i).*passwd.*",
+    r"(?i).*pwd.*",
+    r"(?i).*api_key.*",
+    r"(?i).*apikey.*",
+    r"(?i).*auth.*",
+    r"(?i).*credentials.*",
+    r"(?i).*privatekey.*",
+    r"(?i).*private_key.*",
+    r"(?i).*token.*",
+]
+
+DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS = [r"^__.*"]
+
 LogLevelStr = Literal["fatal", "critical", "error", "warning", "info", "debug"]
 
 Event = TypedDict(
@@ -884,3 +901,129 @@ def strip_string(value, max_length=None):
             "rem": [["!limit", "x", max_length - 3, max_length]],
         },
     )
+
+
+def _compile_patterns(patterns):
+    compiled = []
+    for pattern in patterns:
+        try:
+            compiled.append(re.compile(pattern))
+        except re.error:
+            pass
+    return compiled
+
+
+def _should_ignore_variable(name, ignore_patterns):
+    for pattern in ignore_patterns:
+        if pattern.search(name):
+            return True
+    return False
+
+
+def _should_mask_variable(name, mask_patterns):
+    for pattern in mask_patterns:
+        if pattern.search(name):
+            return True
+    return False
+
+
+def _serialize_variable_value(value, max_length=1024):
+    try:
+        if value is None:
+            return "None"
+        elif isinstance(value, bool):
+            return str(value)
+        elif isinstance(value, (int, float)):
+            result = str(value)
+        elif isinstance(value, str):
+            result = repr(value)
+        else:
+            result = repr(value)
+        
+        # Truncate if too long
+        if len(result) > max_length:
+            result = result[:max_length - 3] + "..."
+        
+        return result
+    except Exception:
+        try:
+            return f"<{type(value).__name__} at {hex(id(value))}>"
+        except Exception:
+            return "<unserializable object>"
+
+
+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):
+    if mask_patterns is None:
+        mask_patterns = []
+    if ignore_patterns is None:
+        ignore_patterns = []
+    
+    compiled_mask = _compile_patterns(mask_patterns)
+    compiled_ignore = _compile_patterns(ignore_patterns)
+    
+    try:
+        local_vars = frame.f_locals.copy()
+    except Exception:
+        return {}
+    
+    simple_vars = {}
+    complex_vars = {}
+    
+    for name, value in local_vars.items():
+        if _should_ignore_variable(name, compiled_ignore):
+            continue
+        
+        if _is_simple_type(value):
+            simple_vars[name] = value
+        else:
+            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 _should_mask_variable(name, compiled_mask):
+            result[name] = "***"
+        else:
+            result[name] = _serialize_variable_value(value, max_length)
+        
+        count += 1
+    
+    return result
+
+
+def extract_code_variables_from_exception(exception_info, mask_patterns, ignore_patterns):
+    exc_type, exc_value, tb = exception_info
+    
+    if tb is None:
+        return {}
+    
+    innermost_tb = None
+    for tb_item in iter_stacks(tb):
+        innermost_tb = tb_item
+    
+    if innermost_tb is None:
+        return {}
+    
+    # Extract variables from the kaboom frame only
+    variables = serialize_code_variables(
+        innermost_tb.tb_frame,
+        mask_patterns=mask_patterns,
+        ignore_patterns=ignore_patterns,
+        max_vars=20,
+        max_length=1024
+    )
+    
+    return variables