auto-capture in with blocks

Author: oliverb123

example.py

@@ -105,26 +105,21 @@
 
 
 # You can add tags to a context, and these are automatically added to exceptions captured within that context.
-# Recommended: you can use the `@posthog.scoped` decorator to enter a new 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
+with posthog.new_context():
+    posthog.tag("user_id", 123)
+    posthog.tag("transaction_id", "abc123")
+    posthog.tag("some_arbitrary_value", {"tags": "can be dicts"})
+    # This exception will be captured with the tags set above
+    raise Exception("Order processing failed")
+
+
+# You can also use the `@posthog.scoped` decorator to enter a new 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")
 
-
-# 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/scopes.py

@@ -1,6 +1,7 @@
 import contextvars
 from contextlib import contextmanager
 from typing import Any, Callable, Dict, TypeVar, cast
+from posthog import capture_exception
 
 _context_stack: contextvars.ContextVar[list] = contextvars.ContextVar("posthog_context_stack", default=[{}])
 
@@ -15,30 +16,24 @@ def new_context():
     # 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 be 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.
+    Any tags set within this scope will be isolated to this context. Any exceptions raised
+    within the context will be captured and tagged with the context tags.
 
     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
+            # The exception will be captured and tagged with the context tags
+            raise ValueError("Something went wrong")
+
     """
     current_stack = _context_stack.get()
     new_stack = current_stack + [{}]
     token = _context_stack.set(new_stack)
     try:
         yield
+    except Exception as e:
+        capture_exception(e)
+        raise
     finally:
         _context_stack.reset(token)
 
@@ -80,8 +75,7 @@ def scoped(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.
+    tags before re-raising it.
 
     Args:
         func: The function to wrap
@@ -96,17 +90,10 @@ def process_payment(payment_id):
     """
     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 func(*args, **kwargs)
+
 
     return cast(F, wrapper)

posthog/test/test_scopes.py

@@ -105,3 +105,30 @@ def failing_function():
 
         # Context should be cleared after function execution
         self.assertEqual(get_tags(), {})
+
+    @patch("posthog.capture_exception")
+    def test_new_context_exception_handling(self, mock_capture):
+        test_exception = RuntimeError("Context exception")
+
+        # Set up outer context
+        tag("outer_context", "outer_value")
+
+        try:
+            with new_context():
+                tag("inner_context", "inner_value")
+                raise test_exception
+        except RuntimeError:
+            pass  # Expected exception
+
+        # Exception should be captured with inner 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 inner context was included in properties
+        self.assertEqual(kwargs.get("properties", {}).get("inner_context"), "inner_value")
+
+        # Outer context should still be intact
+        self.assertEqual(get_tags()["outer_context"], "outer_value")