diff --git a/raven/base.py b/raven/base.py index f258aa2b..74daf94a 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 78252922..258910c6 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 1b1851ac..3cf290b3 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')