feat(err): add context manager and tag functions (#239)

* add context maanager and tag functions

* Update posthog/scopes.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* ran black

* black locally disagrees with ci. python is awful

* bleh

* isort

* fix mypyp thing

* Revert "fix mypyp thing"

This reverts commit 21ad873.

* update baseline

* lets try again

* alright lets try again

* revert to baseline

* try ignoring it i guess

* black

* try supporting async too

* black

* mypy

* formatting

* fix changelog

* fix comment

* we only support python 3.9+

* change decorator name

* fix example

* isort

* auto-capture in with blocks

* fix tests

* black

* inherit tags by default

* assert swap

* add tags to all events

* rm comment

* Update example.py

Co-authored-by: David Newell <[email protected]>

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: David Newell <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -1,3 +1,12 @@
+## 4.3.0
+
+Add exception context management:
+ - New context manager with `posthog.new_context()`
+ - Tag functions: `posthog.tag()`, `posthog.get_tags()`, `posthog.clear_tags()`
+ - Function decorator:
+   - `@posthog.scoped` - Creates context and captures exceptions thrown within the function
+ - Automatic deduplication of exceptions to ensure each exception is only captured once
+
 ## 4.2.1 - 2025-6-05
 
 1. fix: feature flag request use geoip_disable (#235)

example.py

@@ -104,4 +104,44 @@
 print(posthog.get_remote_config_payload("encrypted_payload_flag_key"))
 
 
+# You can add tags to a context, and these are automatically added to any events (including exceptions) captured
+# within that context.
+
+# You can enter a new context using a with statement. Any exceptions thrown in the context will be captured,
+# and tagged with the context tags. Other events captured will also be tagged with the context tags. By default,
+# the new context inherits tags from the parent context.
+with posthog.new_context():
+    posthog.tag("transaction_id", "abc123")
+    posthog.tag("some_arbitrary_value", {"tags": "can be dicts"})
+
+    # This event will be captured with the tags set above
+    posthog.capture("order_processed")
+    # This exception will be captured with the tags set above
+    raise Exception("Order processing failed")
+
+
+# Use fresh=True to start with a clean context (no inherited tags)
+with posthog.new_context(fresh=True):
+    posthog.tag("session_id", "xyz789")
+    # Only session_id tag will be present, no inherited tags
+    raise Exception("Session handling failed")
+
+
+# You can also use the `@posthog.scoped()` decorator to enter a new context.
+# By default, it inherits tags from the parent context
+@posthog.scoped()
+def process_order(order_id):
+    posthog.tag("order_id", order_id)
+    # Exception will be captured and tagged automatically
+    raise Exception("Order processing failed")
+
+
+# Use fresh=True to start with a clean context (no inherited tags)
+@posthog.scoped(fresh=True)
+def process_payment(payment_id):
+    posthog.tag("payment_id", payment_id)
+    # Only payment_id tag will be present, no inherited tags
+    raise Exception("Payment processing failed")
+
+
 posthog.shutdown()

posthog/__init__.py

@@ -4,11 +4,19 @@
 
 from posthog.client import Client
 from posthog.exception_capture import Integrations  # noqa: F401
+from posthog.scopes import clear_tags, get_tags, new_context, scoped, tag
 from posthog.types import FeatureFlag, FlagsAndPayloads
 from posthog.version import VERSION
 
 __version__ = VERSION
 
+"""Context management."""
+new_context = new_context
+tag = tag
+get_tags = get_tags
+clear_tags = clear_tags
+tracked = scoped
+
 """Settings."""
 api_key = None  # type: Optional[str]
 host = None  # type: Optional[str]

posthog/client.py

@@ -29,6 +29,7 @@
     get,
     remote_config,
 )
+from posthog.scopes import get_tags
 from posthog.types import (
     FeatureFlag,
     FeatureFlagResult,
@@ -466,6 +467,11 @@ def capture(
         require("properties", properties, dict)
         require("event", event, string_types)
 
+        # Grab current context tags, if any exist
+        context_tags = get_tags()
+        if context_tags:
+            properties.update(context_tags)
+
         msg = {
             "properties": properties,
             "timestamp": timestamp,
@@ -662,6 +668,11 @@ 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
+
             # 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/exception_integrations/django.py

@@ -32,7 +32,6 @@ class DjangoIntegration:
     identifier = "django"
 
     def __init__(self, capture_exception_fn=None):
-
         if DJANGO_VERSION < (4, 2):
             raise IntegrationEnablingError("Django 4.2 or newer is required.")
 
@@ -60,7 +59,6 @@ def uninstall(self):
 
 
 class DjangoRequestExtractor:
-
     def __init__(self, request):
         # type: (Any) -> None
         self.request = request

posthog/exception_utils.py

@@ -24,7 +24,6 @@
 
 
 if TYPE_CHECKING:
-
     from types import FrameType, TracebackType
     from typing import (  # noqa: F401
         Any,

posthog/scopes.py

@@ -0,0 +1,120 @@
+import contextvars
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, TypeVar, cast
+
+_context_stack: contextvars.ContextVar[list] = contextvars.ContextVar("posthog_context_stack", default=[{}])
+
+
+def _get_current_context() -> Dict[str, Any]:
+    return _context_stack.get()[-1]
+
+
+@contextmanager
+def new_context(fresh=False):
+    """
+     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. Any exceptions raised
+    or events captured within the context will be tagged with the context tags.
+
+     Args:
+         fresh: Whether to start with a fresh context (default: False).
+                If False, inherits tags from parent context.
+                If True, starts with no tags.
+
+     Examples:
+         # Inherit parent context tags
+         with posthog.new_context():
+             posthog.tag("request_id", "123")
+             # Both this event and the exception will be tagged with the context tags
+             posthog.capture("event_name", {"property": "value"})
+             raise ValueError("Something went wrong")
+
+         # Start with fresh context (no inherited tags)
+         with posthog.new_context(fresh=True):
+             posthog.tag("request_id", "123")
+             # Both this event and the exception will be tagged with the context tags
+             posthog.capture("event_name", {"property": "value"})
+             raise ValueError("Something went wrong")
+
+    """
+    import posthog
+
+    current_tags = _get_current_context().copy()
+    current_stack = _context_stack.get()
+    new_stack = current_stack + [{}] if fresh else current_stack + [current_tags]
+    token = _context_stack.set(new_stack)
+
+    try:
+        yield
+    except Exception as e:
+        posthog.capture_exception(e)
+        raise
+    finally:
+        _context_stack.reset(token)
+
+
+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 scoped(fresh=False):
+    """
+    Decorator that creates a new context for the function. Simply wraps
+    the function in a with posthog.new_context(): block.
+
+    Args:
+        fresh: Whether to start with a fresh context (default: False)
+
+    Example:
+        @posthog.scoped()
+        def process_payment(payment_id):
+            posthog.tag("payment_id", payment_id)
+            posthog.tag("payment_method", "credit_card")
+
+            # This event will be captured with tags
+            posthog.capture("payment_started")
+            # If this raises an exception, it will be captured with tags
+            # and then re-raised
+            some_risky_function()
+    """
+
+    def decorator(func: F) -> F:
+        from functools import wraps
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with new_context(fresh=fresh):
+                return func(*args, **kwargs)
+
+        return cast(F, wrapper)
+
+    return decorator