Welcome, guest | Sign In | My Account | Store | Cart

Create objects that act as both context managers and as decorators, and behave the same in both cases.

Works with Python 2.4 - 2.7 and Python 3. The tests require unittest2 or Python 3.2 to run. (And because the tests use the with statement they won't work with Python 2.4.)

Example:

from contextdecorator import ContextDecorator

class mycontext(ContextDecorator):

    def __init__(self, *args):
        """Normal initialiser"""

    def before(self):
        """
        Called on entering the with block or starting the decorated function.

        If used in a with statement whatever this method returns will be the
        context manager.
        """

    def after(self, *exc):
        """
        Called on exit. Arguments and return value of this method have
        the same meaning as the __exit__ method of a normal context
        manager.
        """

Both before and after methods are optional (but providing neither is somewhat pointless). See the tests for more usage examples.

Python, 199 lines
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
# (c) Michael Foord, 2010
# http://voidspace.org.uk/blog
'''
Create objects that act as both context managers *and* as decorators, and behave the same in both cases.

Works with Python 2.4 - 2.7 and Python 3. The tests require unittest2 or Python 3.2 to run.

Example:

from contextdecorator import ContextDecorator

class mycontext(ContextDecorator):

    def __init__(self, *args):
        """Normal initialiser"""

    def before(self):
        """
        Called on entering the with block or starting the decorated function.
        
        If used in a with statement whatever this method returns will be the
        context manager.
        """
    
    def after(self, *exc):
        """
        Called on exit. Arguments and return value of this method have
        the same meaning as the __exit__ method of a normal context
        manager.
        """

@mycontext('some', 'args')
def function():
    pass

with mycontext('some', 'args') as something:
    pass

See the tests for more usage examples.
'''

# Only needed for tests
from __future__ import with_statement

import sys

try:
    from functools import wraps
except ImportError:
    # Python 2.4 compatibility
    def wraps(original):
        def inner(f):
            f.__name__ = original.__name__
            return f
        return inner

# horrible reraise code for compatibility
# with Python 2 & 3
if sys.version_info >= (3,0):
    exec ("""
def _reraise(cls, val, tb):
    raise val
""")
else:
    exec ("""
def _reraise(cls, val, tb):
    raise cls, val, tb
""")


EXC = (None, None, None)

class ContextDecorator(object):
    before = None
    after = None
        
    def __call__(self, f):
        @wraps(f)
        def inner(*args, **kw):
            if self.before is not None:
                self.before()
            
            exc = EXC
            try:
                result = f(*args, **kw)
            except Exception:
                exc = sys.exc_info()
            
            catch = False
            if self.after is not None:
                catch = self.after(*exc)
            
            if not catch and exc is not EXC:
                _reraise(*exc)
            return result
        return inner
            
    def __enter__(self):
        if self.before is not None:
            return self.before()
    
    def __exit__(self, *exc):
        catch = False
        if self.after is not None:
            catch = self.after(*exc)
        return catch


if __name__ == '__main__':
    import sys

    if sys.version_info >= (3, 2):
        import unittest as unittest2
    else:
        import unittest2
    
    class mycontext(ContextDecorator):
        started = False
        exc = None
        catch = False
    
        def before(self):
            self.started = True
            return self
    
        def after(self, *exc):
            self.exc = exc
            return self.catch


    class TestContext(unittest2.TestCase):

        def test_context(self):
            context = mycontext()
            with context as result:
                self.assertIs(result, context)
                self.assertTrue(context.started)
            
            self.assertEqual(context.exc, (None, None, None))
    
        def test_context_with_exception(self):
            context = mycontext()
        
            with self.assertRaisesRegexp(NameError, 'foo'):
                with context:
                    raise NameError('foo')
        
            context.exc = (None, None, None)
            context.catch = True
            with context:
                raise NameError('foo')
            self.assertNotEqual(context.exc, (None, None, None))
        
        def test_decorator(self):
            context = mycontext()
        
            @context
            def test():
                self.assertIsNone(context.exc)
                self.assertTrue(context.started)
            test()
            self.assertEqual(context.exc, (None, None, None))
    
        def test_decorator_with_exception(self):
            context = mycontext()
        
            @context
            def test():
                self.assertIsNone(context.exc)
                self.assertTrue(context.started)
                raise NameError('foo')
        
            with self.assertRaisesRegexp(NameError, 'foo'):
                test()
            self.assertNotEqual(context.exc, (None, None, None))

        def test_decorating_method(self):
            context = mycontext()
        
            class Test(object):
            
                @context
                def method(self, a, b, c=None):
                    self.a = a
                    self.b = b
                    self.c = c
        
            test = Test()
            test.method(1, 2)
            self.assertEqual(test.a, 1)
            self.assertEqual(test.b, 2)
            self.assertEqual(test.c, None)

            test.method('a', 'b', 'c')
            self.assertEqual(test.a, 'a')
            self.assertEqual(test.b, 'b')
            self.assertEqual(test.c, 'c')

    unittest2.main()

5 comments

Michael Foord (author) 11 years, 6 months ago  # | flag

I wrote this after writing almost identical code the second time for "patch" in the mock module. (The patch decorator can be used as a decorator or as a context manager and I was writing a new variant.) It turns out that both py.test and django have similar code in places - so it is not an uncommon pattern.

Douglas Napoleone 11 years, 6 months ago  # | flag

Fantastic! I wrote something similar, but not to this level or with this version support. Only request: kwdarg support. Absolutely amazing. This will save me soo much time going forward.

Michael Foord (author) 11 years, 6 months ago  # | flag

@Doug What do you mean by keyword argument support? As you implement your own __init__ method you can take whatever args you want. If you mean in the start and finish methods could you give me an example of how you would like to use it.

Benjamin Peterson 11 years, 6 months ago  # | flag

You don't need that "assert" and __tracebackhide__ crud in your reraise hack. :)

Michael Foord (author) 11 years, 6 months ago  # | flag

Ha! Thanks Benjamin. I just copied and pasted that hack from someone... ;-)