add context maanager and tag functions

Author: oliverb123

CHANGELOG.md

@@ -1,3 +1,12 @@
+## 4.1.0 - 2025-05-01
+
+1. Add exception context management:
+   - New context manager with `posthog.new_context()`
+   - Tag functions: `posthog.tag()`, `posthog.get_tags()`, `posthog.clear_tags()`
+   - Function decorator:
+     - `@posthog.tracked` - Creates context and captures exceptions thrown within the function
+   - Automatic deduplication of exceptions to ensure each exception is only captured once
+
 ## 4.0.1 – 2025-04-29
 
 1. Remove deprecated `monotonic` library. Use Python's core `time.monotonic` function instead

example.py

@@ -104,4 +104,26 @@
 print(posthog.get_remote_config_payload("encrypted_payload_flag_key"))
 
 
+# You can add tags to a context, and these are automatically added to exceptions captured within that context.
+# Recommended: you can use the `@posthog.tracked` decorator.
+@posthog.tracked
+def process_order(order_id):
+    posthog.tag("order_id", order_id)
+    # Exception will be captured and tagged automatically
+    raise Exception("Order processing failed")
+
+# You can also use the directly context manager, but you should prefer to use the decorator where possible,
+# because it's tricky to remember to always capture the exception within the context scope.
+with posthog.new_context():
+    posthog.tag("user_id", "123")
+    posthog.tag("transaction_id", "abc123")
+    try:
+        # Any exception within this block will include these tags
+        raise Exception("Order processing failed")
+    except Exception as e:
+        # Note the exception must be captured within the scope for the tags
+        # to be applied
+        posthog.capture_exception(e)
+
+
 posthog.shutdown()

posthog/__init__.py

@@ -5,10 +5,18 @@
 from posthog.client import Client
 from posthog.exception_capture import Integrations  # noqa: F401
 from posthog.types import FeatureFlag, FlagsAndPayloads
+from posthog.scopes import new_context, tag, get_tags, clear_tags, tracked
 from posthog.version import VERSION
 
 __version__ = VERSION
 
+"""Expose context management functions at module level."""
+new_context = new_context
+tag = tag
+get_tags = get_tags
+clear_tags = clear_tags
+tracked = tracked
+
 """Settings."""
 api_key = None  # type: Optional[str]
 host = None  # type: Optional[str]

posthog/client.py

@@ -14,6 +14,7 @@
 from dateutil.tz import tzutc
 from six import string_types
 
+from posthog.scopes import get_tags
 from posthog.consumer import Consumer
 from posthog.exception_capture import ExceptionCapture
 from posthog.exception_utils import exc_info_from_error, exceptions_from_error_tuple, handle_in_app
@@ -662,6 +663,16 @@ def capture_exception(
         try:
             properties = properties or {}
 
+            # Check if this exception has already been captured
+            if exception is not None and hasattr(exception, "__posthog_exception_captured"):
+                self.log.debug("Exception already captured, skipping")
+                return
+
+            # Grab current context tags, if any exist
+            context_tags = get_tags()
+            if context_tags:
+                properties.update(context_tags)
+
             # if there's no distinct_id, we'll generate one and set personless mode
             # via $process_person_profile = false
             if distinct_id is None:
@@ -705,7 +716,13 @@ def capture_exception(
             if self.log_captured_exceptions:
                 self.log.exception(exception, extra=kwargs)
 
-            return self.capture(distinct_id, "$exception", properties, context, timestamp, uuid, groups)
+            res = self.capture(distinct_id, "$exception", properties, context, timestamp, uuid, groups)
+
+            # Mark the exception as captured to prevent duplicate captures
+            if exception is not None:
+                setattr(exception, "__posthog_exception_captured", True)
+
+            return res
         except Exception as e:
             self.log.exception(f"Failed to capture exception: {e}")
 

posthog/scopes.py

@@ -0,0 +1,114 @@
+import threading
+from contextlib import contextmanager
+from typing import Any, Dict, Callable, TypeVar, cast
+
+_scopes_local = threading.local()
+
+def _init_guard() -> None:
+    if not hasattr(_scopes_local, "context_stack"):
+        _scopes_local.context_stack = [{}]
+
+
+def _get_current_context() -> Dict[str, Any]:
+    _init_guard()
+    return _scopes_local.context_stack[-1]
+
+
+@contextmanager
+def new_context():
+    # TODO - we could extend this context idea to also apply to other event types eventually,
+    # but right now it only applies to exceptions...
+    """
+    Create a new context scope that will be active for the duration of the with block.
+    Any tags set within this scope will be isolated to this context. Tags added to a
+    context will be added to exceptions captured within that context.
+
+    NOTE: tags set within a context will only added to exceptions captured within that
+    context - ensure you call `posthog.capture_exception()` before the end of the with
+    block, or the extra tags will be lost.
+
+    It's strongly recommended to use the `posthog.tracked` decorator to instrument functions, rather
+    than directly using this context manager.
+
+    Example:
+        with posthog.new_context():
+            posthog.tag("user_id", "123")
+            try:
+                # Do something that might raise an exception
+            except Exception as e:
+                posthog.capture_exception(e)
+                raise e
+    """
+    _init_guard()
+    _scopes_local.context_stack.append({})
+    try:
+        yield
+    finally:
+        if len(_scopes_local.context_stack) > 1:
+            _scopes_local.context_stack.pop()
+
+
+def tag(key: str, value: Any) -> None:
+    """
+    Add a tag to the current context.
+
+    Args:
+        key: The tag key
+        value: The tag value
+
+    Example:
+        posthog.tag("user_id", "123")
+    """
+    _get_current_context()[key] = value
+
+
+def get_tags() -> Dict[str, Any]:
+    """
+    Get all tags from the current context. Note, modifying
+    the returned dictionary will not affect the current context.
+
+    Returns:
+        Dict of all tags in the current context
+    """
+    return _get_current_context().copy()
+
+
+def clear_tags() -> None:
+    """Clear all tags in the current context."""
+    _get_current_context().clear()
+
+
+F = TypeVar('F', bound=Callable[..., Any])
+def tracked(func: F) -> F:
+    """
+    Decorator that creates a new context for the function, wraps the function in a
+    try/except block, and if an exception occurs, captures it with the current context
+    tags before re-raising it. This is the recommended way to wrap/track functions for
+    posthog error tracking.
+
+    Args:
+        func: The function to wrap
+
+    Example:
+        @posthog.tracked
+        def process_payment(payment_id):
+            posthog.tag("payment_id", payment_id)
+            posthog.tag("payment_method", "credit_card")
+            # If this raises an exception, it will be captured with tags
+            # and then re-raised
+    """
+    from functools import wraps
+    import posthog
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        with new_context():
+            try:
+                return func(*args, **kwargs)
+            except Exception as e:
+                # Capture the exception with current context tags
+                # The capture_exception function will handle deduplication
+                posthog.capture_exception(e, properties=get_tags())
+                raise  # Re-raise the exception after capturing it
+
+    return cast(F, wrapper)

posthog/test/test_scopes.py

@@ -0,0 +1,107 @@
+import unittest
+from unittest.mock import patch
+
+from posthog.scopes import new_context, tag, get_tags, clear_tags, tracked
+
+
+class TestScopes(unittest.TestCase):
+    def setUp(self):
+        # Reset any context between tests
+        clear_tags()
+
+    def test_tag_and_get_tags(self):
+        tag("key1", "value1")
+        tag("key2", 2)
+
+        tags = get_tags()
+        self.assertEqual(tags["key1"], "value1")
+        self.assertEqual(tags["key2"], 2)
+
+    def test_clear_tags(self):
+        tag("key1", "value1")
+        self.assertEqual(get_tags()["key1"], "value1")
+
+        clear_tags()
+        self.assertEqual(get_tags(), {})
+
+    def test_new_context_isolation(self):
+        # Set tag in outer context
+        tag("outer", "value")
+
+        with new_context():
+            # Inner context should start empty
+            self.assertEqual(get_tags(), {})
+
+            # Set tag in inner context
+            tag("inner", "value")
+            self.assertEqual(get_tags()["inner"], "value")
+
+            # Outer tag should not be visible
+            self.assertNotIn("outer", get_tags())
+
+        # After exiting context, inner tag should be gone
+        self.assertNotIn("inner", get_tags())
+
+        # Outer tag should still be there
+        self.assertEqual(get_tags()["outer"], "value")
+
+    def test_nested_contexts(self):
+        tag("level1", "value1")
+
+        with new_context():
+            tag("level2", "value2")
+
+            with new_context():
+                tag("level3", "value3")
+                self.assertEqual(get_tags(), {"level3": "value3"})
+
+            # Back to level 2
+            self.assertEqual(get_tags(), {"level2": "value2"})
+
+        # Back to level 1
+        self.assertEqual(get_tags(), {"level1": "value1"})
+
+    @patch('posthog.capture_exception')
+    def test_tracked_decorator_success(self, mock_capture):
+        @tracked
+        def successful_function(x, y):
+            tag("x", x)
+            tag("y", y)
+            return x + y
+
+        result = successful_function(1, 2)
+
+        # Function should execute normally
+        self.assertEqual(result, 3)
+
+        # No exception should be captured
+        mock_capture.assert_not_called()
+
+        # Context should be cleared after function execution
+        self.assertEqual(get_tags(), {})
+
+    @patch('posthog.capture_exception')
+    def test_tracked_decorator_exception(self, mock_capture):
+        test_exception = ValueError("Test exception")
+
+        @tracked
+        def failing_function():
+            tag("important_context", "value")
+            raise test_exception
+
+        # Function should raise the exception
+        with self.assertRaises(ValueError):
+            failing_function()
+
+        # Exception should be captured with context
+        mock_capture.assert_called_once()
+        args, kwargs = mock_capture.call_args
+
+        # Check that the exception was passed
+        self.assertEqual(args[0], test_exception)
+
+        # Check that the context was included in properties
+        self.assertEqual(kwargs.get('properties', {}).get('important_context'), "value")
+
+        # Context should be cleared after function execution
+        self.assertEqual(get_tags(), {})