Allow push_scope to act as a decorator

sentry_sdk/_compat.py

@@ -15,6 +15,35 @@
 
 PY2 = sys.version_info[0] == 2
 
+if not PY2 and sys.version_info >= (3, 2):
+    from contextlib import ContextDecorator
+else:
+    from functools import wraps
+
+    class ContextDecorator(object):
+        "A base class or mixin that enables context managers to work as decorators."
+
+        def _recreate_cm(self):
+            """Return a recreated instance of self.
+
+            Allows an otherwise one-shot context manager like
+            _GeneratorContextManager to support use as
+            a decorator via implicit recreation.
+
+            This is a private interface just for _GeneratorContextManager.
+            See issue #11647 for details.
+            """
+            return self
+
+        def __call__(self, func):
+            @wraps(func)
+            def inner(*args, **kwds):
+                with self._recreate_cm():
+                    return func(*args, **kwds)
+
+            return inner
+
+
 if PY2:
     import urlparse  # noqa
 

sentry_sdk/api.py

@@ -1,13 +1,14 @@
 import inspect
 from contextlib import contextmanager
 
-from sentry_sdk.hub import Hub
+from sentry_sdk.hub import Hub, ScopeManager
 from sentry_sdk.scope import Scope
 
 from sentry_sdk._types import MYPY
 
 if MYPY:
     from typing import Any
+    from typing import Dict
     from typing import Optional
     from typing import overload
     from typing import Callable
@@ -36,6 +37,10 @@ def overload(x):
     "flush",
     "last_event_id",
     "start_span",
+    "set_tag",
+    "set_extra",
+    "set_user",
+    "set_level",
 ]
 
 
@@ -48,6 +53,15 @@ def hubmethod(f):
     return f
 
 
+def scopemethod(f):
+    # type: (F) -> F
+    f.__doc__ = "%s\n\n%s" % (
+        "Alias for :py:meth:`sentry_sdk.Scope.%s`" % f.__name__,
+        inspect.getdoc(getattr(Scope, f.__name__)),
+    )
+    return f
+
+
 @hubmethod
 def capture_event(
     event,  # type: Event
@@ -131,7 +145,7 @@ def inner():
 
 @overload  # noqa
 def push_scope():
-    # type: () -> ContextManager[Scope]
+    # type: () -> ScopeManager
     pass
 
 
@@ -147,7 +161,7 @@ def push_scope(
 def push_scope(
     callback=None  # type: Optional[Callable[[Scope], None]]
 ):
-    # type: (...) -> Optional[ContextManager[Scope]]
+    # type: (...) -> Optional[ScopeManager]
     hub = Hub.current
     if hub is not None:
         return hub.push_scope(callback)
@@ -163,6 +177,38 @@ def inner():
         return None
 
 
+@scopemethod  # noqa
+def set_tag(key, value):
+    # type: (str, Any) -> None
+    hub = Hub.current
+    if hub is not None:
+        hub.current_scope.set_tag(key, value)
+
+
+@scopemethod  # noqa
+def set_extra(key, value):
+    # type: (str, Any) -> None
+    hub = Hub.current
+    if hub is not None:
+        hub.current_scope.set_extra(key, value)
+
+
+@scopemethod  # noqa
+def set_user(value):
+    # type: (Dict[str, Any]) -> None
+    hub = Hub.current
+    if hub is not None:
+        hub.current_scope.set_user(value)
+
+
+@scopemethod  # noqa
+def set_level(value):
+    # type: (str) -> None
+    hub = Hub.current
+    if hub is not None:
+        hub.current_scope.set_level(value)
+
+
 @hubmethod
 def flush(
     timeout=None,  # type: Optional[float]

sentry_sdk/hub.py

@@ -5,7 +5,7 @@
 from datetime import datetime
 from contextlib import contextmanager
 
-from sentry_sdk._compat import with_metaclass
+from sentry_sdk._compat import with_metaclass, ContextDecorator
 from sentry_sdk.scope import Scope
 from sentry_sdk.client import Client
 from sentry_sdk.tracing import Span
@@ -123,15 +123,20 @@ def main(cls):
         return GLOBAL_HUB
 
 
-class _ScopeManager(object):
+class ScopeManager(ContextDecorator):
     def __init__(self, hub):
         # type: (Hub) -> None
         self._hub = hub
-        self._original_len = len(hub._stack)
-        self._layer = hub._stack[-1]
 
     def __enter__(self):
         # type: () -> Scope
+        client, scope = self._hub._stack[-1]
+        new_layer = (client, copy.copy(scope))
+        self._hub._stack.append(new_layer)
+
+        self._original_len = len(self._hub._stack)
+        self._layer = self._hub._stack[-1]
+
         scope = self._layer[1]
         assert scope is not None
         return scope
@@ -263,6 +268,12 @@ def client(self):
         """Returns the current client on the hub."""
         return self._stack[-1][0]
 
+    @property
+    def current_scope(self):
+        # type: () -> Scope
+        """Returns the current scope on the hub."""
+        return self._stack[-1][1]
+
     def last_event_id(self):
         # type: () -> Optional[str]
         """Returns the last event ID."""
@@ -448,7 +459,7 @@ def start_span(
     def push_scope(
         self, callback=None  # type: Optional[None]
     ):
-        # type: (...) -> ContextManager[Scope]
+        # type: (...) -> ScopeManager
         pass
 
     @overload  # noqa
@@ -461,14 +472,14 @@ def push_scope(
     def push_scope(  # noqa
         self, callback=None  # type: Optional[Callable[[Scope], None]]
     ):
-        # type: (...) -> Optional[ContextManager[Scope]]
+        # type: (...) -> Optional[ScopeManager]
         """
         Pushes a new layer on the scope stack.
 
         :param callback: If provided, this method pushes a scope, calls
             `callback`, and pops the scope again.
 
-        :returns: If no `callback` is provided, a context manager that should
+        :returns: If no `callback` is provided, a ContextDecorator that should
             be used to pop the scope again.
         """
 
@@ -477,11 +488,7 @@ def push_scope(  # noqa
                 callback(scope)
             return None
 
-        client, scope = self._stack[-1]
-        new_layer = (client, copy.copy(scope))
-        self._stack.append(new_layer)
-
-        return _ScopeManager(self)
+        return ScopeManager(self)
 
     scope = push_scope
 

sentry_sdk/scope.py

@@ -123,6 +123,11 @@ def level(self, value):
         """When set this overrides the level."""
         self._level = value
 
+    def set_level(self, value):
+        # type: (Optional[str]) -> None
+        """When set this overrides the level."""
+        self._level = value
+
     @_attr_setter
     def fingerprint(self, value):
         # type: (Optional[List[str]]) -> None
@@ -144,6 +149,11 @@ def user(self, value):
         """When set a specific user is bound to the scope."""
         self._user = value
 
+    def set_user(self, value):
+        # type: (Dict[str, Any]) -> None
+        """When set a specific user is bound to the scope."""
+        self._user = value
+
     @property
     def span(self):
         # type: () -> Optional[Span]

tests/test_basics.py

@@ -11,6 +11,7 @@
     add_breadcrumb,
     last_event_id,
     Hub,
+    set_level,
 )
 from sentry_sdk.integrations.logging import LoggingIntegration
 
@@ -174,6 +175,30 @@ def _(scope):
     assert Hub.current._stack[-1][1] is outer_scope
 
 
+def test_push_scope_as_decorator(sentry_init, capture_events):
+    sentry_init()
+    events = capture_events()
+
+    Hub.current.bind_client(None)
+
+    outer_scope = Hub.current._stack[-1][1]
+
+    @push_scope()
+    def foo():
+        set_level("warning")
+        try:
+            1 / 0
+        except Exception as e:
+            capture_exception(e)
+
+    foo()
+
+    event, = events
+    assert outer_scope._level != "warning"
+    assert event["level"] == "warning"
+    assert "exception" in event
+
+
 def test_breadcrumbs(sentry_init, capture_events):
     sentry_init(max_breadcrumbs=10)
     events = capture_events()
@@ -240,23 +265,6 @@ def test_client_initialized_within_scope(sentry_init, caplog):
     assert record.msg.startswith("init() called inside of pushed scope.")
 
 
-def test_scope_leaks_cleaned_up(sentry_init, caplog):
-    caplog.set_level(logging.WARNING)
-
-    sentry_init(debug=True)
-
-    old_stack = list(Hub.current._stack)
-
-    with push_scope():
-        push_scope()
-
-    assert Hub.current._stack == old_stack
-
-    record, = (x for x in caplog.records if x.levelname == "WARNING")
-
-    assert record.message.startswith("Leaked 1 scopes:")
-
-
 def test_scope_popped_too_soon(sentry_init, caplog):
     caplog.set_level(logging.ERROR)