feat: Added support for with-scope like behavior (getsentry#65)

mitsuhiko · web-flow · commit c99d54d0955b · 2018-09-17T22:03:19.000+02:00
diff --git a/sentry_sdk/api.py b/sentry_sdk/api.py
@@ -105,6 +105,20 @@ def inner():
         return inner()
 
 
+@hubmethod
+def push_scope(callback=None):
+    hub = Hub.current
+    if hub is not None:
+        return hub.push_scope(callback)
+    elif callback is None:
+
+        @contextmanager
+        def inner():
+            yield Scope()
+
+        return inner()
+
+
 @hubmethod
 def last_event_id():
     hub = Hub.current
diff --git a/sentry_sdk/hub.py b/sentry_sdk/hub.py
@@ -60,7 +60,10 @@ def __init__(self, hub, layer):
         self._layer = layer
 
     def __enter__(self):
-        return self
+        scope = self._layer[1]
+        if scope is None:
+            scope = Scope()
+        return scope
 
     def __exit__(self, exc_type, exc_value, tb):
         assert self._hub.pop_scope_unsafe() == self._layer, "popped wrong scope"
@@ -205,13 +208,20 @@ def add_breadcrumb(self, crumb=None, hint=None, **kwargs):
         while len(scope._breadcrumbs) >= client.options["max_breadcrumbs"]:
             scope._breadcrumbs.popleft()
 
-    def push_scope(self):
+    def push_scope(self, callback=None):
         """Pushes a new layer on the scope stack. Returns a context manager
-        that should be used to pop the scope again."""
+        that should be used to pop the scope again.  Alternatively a callback
+        can be provided that is executed in the context of the scope.
+        """
         client, scope = self._stack[-1]
         new_layer = (client, copy.copy(scope))
         self._stack.append(new_layer)
-        return _ScopeManager(self, new_layer)
+
+        if callback is not None:
+            if client is not None:
+                callback(scope)
+        else:
+            return _ScopeManager(self, new_layer)
 
     def pop_scope_unsafe(self):
         """Pops a scope layer from the stack. Try to use the context manager
@@ -224,18 +234,28 @@ def configure_scope(self, callback=None):
         """Reconfigures the scope."""
         client, scope = self._stack[-1]
         if callback is not None:
-            if client is not None and scope is not None:
+            if client is not None:
                 callback(scope)
             return
 
         @contextmanager
         def inner():
-            if client is not None and scope is not None:
+            if client is not None:
                 yield scope
             else:
                 yield Scope()
 
         return inner()
 
+    def scope(self, callback=None):
+        """Pushes a new scope and yields it for configuration.
+
+        The scope is dropped at the end of the with statement.  Alternatively
+        a callback can be provided similar to `configure_scope`.
+        """
+        with self.push_scope():
+            client, scope = self._stack[-1]
+            return self.configure_scope(callback)
+
 
 GLOBAL_HUB = Hub()
diff --git a/sentry_sdk/scope.py b/sentry_sdk/scope.py
@@ -11,6 +11,7 @@ class Scope(object):
     """
 
     __slots__ = (
+        "_level",
         "_fingerprint",
         "_transaction",
         "_user",
@@ -28,6 +29,11 @@ def __init__(self):
 
         self.clear()
 
+    @_attr_setter
+    def level(self, value):
+        """When set this overrides the level."""
+        self._level = value
+
     @_attr_setter
     def fingerprint(self, value):
         """When set this overrides the default fingerprint."""
@@ -69,6 +75,7 @@ def remove_extra(self, key):
 
     def clear(self):
         """Clears the entire scope."""
+        self._level = None
         self._fingerprint = None
         self._transaction = None
         self._user = None
@@ -112,6 +119,9 @@ def apply_to_event(self, event, hint=None):
         def _drop(event, cause, ty):
             logger.info("%s (%s) dropped event (%s)", ty, cause, event)
 
+        if self._level is not None:
+            event["level"] = self._level
+
         event.setdefault("breadcrumbs", []).extend(self._breadcrumbs)
         if event.get("user") is None and self._user is not None:
             event["user"] = self._user
@@ -151,6 +161,7 @@ def _drop(event, cause, ty):
     def __copy__(self):
         rv = object.__new__(self.__class__)
 
+        rv._level = self._level
         rv._fingerprint = self._fingerprint
         rv._transaction = self._transaction
         rv._user = self._user
diff --git a/tests/test_basics.py b/tests/test_basics.py
@@ -1,4 +1,5 @@
 from sentry_sdk import (
+    push_scope,
     configure_scope,
     capture_exception,
     add_breadcrumb,
@@ -105,3 +106,36 @@ def before_breadcrumb(crumb, hint):
     assert_hint.clear()
     add_breadcrumb(foo=42)
     add_breadcrumb(crumb=dict(foo=42))
+
+
+def test_push_scope(sentry_init, capture_events):
+    sentry_init()
+    events = capture_events()
+
+    with push_scope() as scope:
+        scope.level = "warning"
+        try:
+            1 / 0
+        except Exception as e:
+            capture_exception(e)
+
+    event, = events
+
+    assert event["level"] == "warning"
+    assert "exception" in event
+
+
+def test_push_scope_null_client(sentry_init, capture_events):
+    sentry_init()
+    events = capture_events()
+
+    Hub.current.bind_client(None)
+
+    with push_scope() as scope:
+        scope.level = "warning"
+        try:
+            1 / 0
+        except Exception as e:
+            capture_exception(e)
+
+    assert len(events) == 0
