Modify client.capture_exceptions to return both decorator and context manager

3 files changed, 82 insertions, 20 deletions

diff --git a/raven/base.py b/raven/base.py
index f258aa2..74daf94 100644
--- a/raven/base.py
+++ b/raven/base.py
@@ -17,10 +17,14 @@ import uuid
 import warnings
 
 from datetime import datetime
-from functools import wraps
 from pprint import pformat
 from types import FunctionType
 
+if sys.version_info >= (3, 2):
+    import contextlib
+else:
+    import contextlib2 as contextlib
+
 import raven
 from raven.conf import defaults
 from raven.conf.remote import RemoteConfig
@@ -689,10 +693,11 @@ class Client(object):
         return self.capture(
             'raven.events.Exception', exc_info=exc_info, **kwargs)
 
-    def capture_exceptions(self, function_or_exceptions, **kwargs):
+    def capture_exceptions(self, function_or_exceptions=None, **kwargs):
         """
-        Wrap a function in try/except and automatically call ``.captureException``
-        if it raises an exception, then the exception is reraised.
+        Wrap a function or code block in try/except and automatically call
+        ``.captureException`` if it raises an exception, then the exception
+        is reraised.
 
         By default, it will capture ``Exception``
 
@@ -700,28 +705,41 @@ class Client(object):
         >>> def foo():
         >>>     raise Exception()
 
+        >>> with client.capture_exceptions():
+        >>>    raise Exception()
+
         You can also specify exceptions to be caught specifically
 
         >>> @client.capture_exceptions((IOError, LookupError))
         >>> def bar():
         >>>     ...
 
+        >>> with client.capture_exceptions((IOError, LookupError)):
+        >>>    ...
+
         ``kwargs`` are passed through to ``.captureException``.
         """
-        def make_decorator(exceptions):
-            def decorator(func):
-                @wraps(func)
-                def wrapper(*funcargs, **funckwargs):
-                    try:
-                        return func(*funcargs, **funckwargs)
-                    except exceptions:
-                        self.captureException(**kwargs)
-                        raise
-                return wrapper
-            return decorator
+        function = None
+        exceptions = (Exception,)
         if isinstance(function_or_exceptions, FunctionType):
-            return make_decorator((Exception,))(function_or_exceptions)
-        return make_decorator(function_or_exceptions)
+            function = function_or_exceptions
+        elif function_or_exceptions is not None:
+            exceptions = function_or_exceptions
+
+        # In python3.2 contextmanager acts both as contextmanager and decorator
+        @contextlib.contextmanager
+        def make_decorator(exceptions):
+            try:
+                yield
+            except exceptions:
+                self.captureException(**kwargs)
+                raise
+
+        decorator = make_decorator(exceptions)
+
+        if function:
+            return decorator(function)
+        return decorator
 
     def captureQuery(self, query, params=(), engine=None, **kwargs):
         """
diff --git a/setup.py b/setup.py
index 7825292..258910c 100755
--- a/setup.py
+++ b/setup.py
@@ -24,6 +24,10 @@ from setuptools import setup, find_packages
 from setuptools.command.test import test as TestCommand
 import sys
 
+install_requires = [
+    'contextlib2',
+]
+
 setup_requires = [
     'pytest',
 ]
@@ -52,6 +56,10 @@ if sys.version_info[0] == 3:
     unittest2_requires = []
     webpy_tests_requires = []
 
+    # If it's python3.2 or greater, don't use contextlib backport
+    if sys.version_info[1] >= 2:
+        install_requires.remove('contextlib2')
+
 tests_require = [
     'bottle',
     'celery>=2.5',
@@ -110,6 +118,7 @@ setup(
     },
     license='BSD',
     tests_require=tests_require,
+    install_requires=install_requires,
     cmdclass={'test': PyTest},
     include_package_data=True,
     entry_points={
diff --git a/tests/base/tests.py b/tests/base/tests.py
index 1b1851a..3cf290b 100644
--- a/tests/base/tests.py
+++ b/tests/base/tests.py
@@ -323,9 +323,9 @@ class ClientTest(TestCase):
         self.assertEquals(exc['type'], 'DecoratorTestException')
         self.assertEquals(exc['module'], self.DecoratorTestException.__module__)
         stacktrace = exc['stacktrace']
-        # this is a wrapped function so two frames are expected
-        self.assertEquals(len(stacktrace['frames']), 2)
-        frame = stacktrace['frames'][1]
+        # this is a wrapped class object with __call__ so three frames are expected
+        self.assertEquals(len(stacktrace['frames']), 3)
+        frame = stacktrace['frames'][-1]
         self.assertEquals(frame['module'], __name__)
         self.assertEquals(frame['function'], 'test2')
 
@@ -341,6 +341,41 @@ class ClientTest(TestCase):
 
         self.assertEquals(len(self.client.events), 0)
 
+    def test_context_manager_functionality(self):
+        def test4():
+            raise self.DecoratorTestException()
+
+        try:
+            with self.client.capture_exceptions():
+                test4()
+        except self.DecoratorTestException:
+            pass
+
+        self.assertEquals(len(self.client.events), 1)
+        event = self.client.events.pop(0)
+        self.assertEquals(event['message'], 'DecoratorTestException')
+        exc = event['exception']['values'][0]
+        self.assertEquals(exc['type'], 'DecoratorTestException')
+        self.assertEquals(exc['module'], self.DecoratorTestException.__module__)
+        stacktrace = exc['stacktrace']
+        # three frames are expected: test4, `with` block and context manager internals
+        self.assertEquals(len(stacktrace['frames']), 3)
+        frame = stacktrace['frames'][-1]
+        self.assertEquals(frame['module'], __name__)
+        self.assertEquals(frame['function'], 'test4')
+
+    def test_content_manager_filtering(self):
+        def test5():
+            raise Exception()
+
+        try:
+            with self.client.capture_exceptions(self.DecoratorTestException):
+                test5()
+        except Exception:
+            pass
+
+        self.assertEquals(len(self.client.events), 0)
+
     def test_message_event(self):
         self.client.captureMessage(message='test')