Issue 9110. Adding ContextDecorator to contextlib. This enables the creation of APIs that act as decorators as well as context managers. contextlib.contextmanager changed to use ContextDecorator.

4 files changed, 237 insertions, 2 deletions

diff --git a/Doc/library/contextlib.rst b/Doc/library/contextlib.rst
index 2ee9e8d..7a46834 100644
--- a/Doc/library/contextlib.rst
+++ b/Doc/library/contextlib.rst
@@ -51,6 +51,11 @@ Functions provided:
    the exception has been handled, and execution will resume with the statement
    immediately following the :keyword:`with` statement.
 
+   contextmanager uses :class:`ContextDecorator` so the context managers it
+   creates can be used as decorators as well as in :keyword:`with` statements.
+
+   .. versionchanged:: 3.2
+      Use of :class:`ContextDecorator`.
 
 .. function:: closing(thing)
 
@@ -79,6 +84,58 @@ Functions provided:
    ``page.close()`` will be called when the :keyword:`with` block is exited.
 
 
+.. class:: ContextDecorator()
+
+   A base class that enables a context manager to also be used as a decorator.
+
+   Context managers inheriting from ``ContextDecorator`` have to implement
+   ``__enter__`` and ``__exit__`` as normal. ``__exit__`` retains its optional
+   exception handling even when used as a decorator.
+
+   Example::
+
+      from contextlib import ContextDecorator
+
+      class mycontext(ContextDecorator):
+         def __enter__(self):
+            print('Starting')
+            return self
+
+         def __exit__(self, *exc):
+            print('Finishing')
+            return False
+
+      >>> @mycontext()
+      ... def function():
+      ...    print('The bit in the middle')
+      ...
+      >>> function()
+      Starting
+      The bit in the middle
+      Finishing
+
+      >>> with mycontext():
+      ...    print('The bit in the middle')
+      ...
+      Starting
+      The bit in the middle
+      Finishing
+
+   Existing context managers that already have a base class can be extended by
+   using ``ContextDecorator`` as a mixin class::
+
+      from contextlib import ContextDecorator
+
+      class mycontext(ContextBaseClass, ContextDecorator):
+         def __enter__(self):
+            return self
+
+         def __exit__(self, *exc):
+            return False
+
+   .. versionadded:: 3.2
+
+
 .. seealso::
 
    :pep:`0343` - The "with" statement
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index e26d77a..e37fde8 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -4,9 +4,20 @@ import sys
 from functools import wraps
 from warnings import warn
 
-__all__ = ["contextmanager", "closing"]
+__all__ = ["contextmanager", "closing", "ContextDecorator"]
 
-class GeneratorContextManager(object):
+
+class ContextDecorator(object):
+    "A base class or mixin that enables context managers to work as decorators."
+    def __call__(self, func):
+        @wraps(func)
+        def inner(*args, **kwds):
+            with self:
+                return func(*args, **kwds)
+        return inner
+
+
+class GeneratorContextManager(ContextDecorator):
     """Helper for @contextmanager decorator."""
 
     def __init__(self, gen):
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
index 389e7d6..a3e9b07 100644
--- a/Lib/test/test_contextlib.py
+++ b/Lib/test/test_contextlib.py
@@ -202,6 +202,169 @@ class LockContextTestCase(unittest.TestCase):
                 return True
         self.boilerPlate(lock, locked)
 
+
+class mycontext(ContextDecorator):
+    started = False
+    exc = None
+    catch = False
+
+    def __enter__(self):
+        self.started = True
+        return self
+
+    def __exit__(self, *exc):
+        self.exc = exc
+        return self.catch
+
+
+class TestContextDecorator(unittest.TestCase):
+
+    def test_contextdecorator(self):
+        context = mycontext()
+        with context as result:
+            self.assertIs(result, context)
+            self.assertTrue(context.started)
+
+        self.assertEqual(context.exc, (None, None, None))
+
+
+    def test_contextdecorator_with_exception(self):
+        context = mycontext()
+
+        with self.assertRaisesRegexp(NameError, 'foo'):
+            with context:
+                raise NameError('foo')
+        self.assertIsNotNone(context.exc)
+        self.assertIs(context.exc[0], NameError)
+
+        context = mycontext()
+        context.catch = True
+        with context:
+            raise NameError('foo')
+        self.assertIsNotNone(context.exc)
+        self.assertIs(context.exc[0], NameError)
+
+
+    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.assertIsNotNone(context.exc)
+        self.assertIs(context.exc[0], NameError)
+
+
+    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
+
+        # these tests are for argument passing when used as a decorator
+        test = Test()
+        test.method(1, 2)
+        self.assertEqual(test.a, 1)
+        self.assertEqual(test.b, 2)
+        self.assertEqual(test.c, None)
+
+        test = Test()
+        test.method('a', 'b', 'c')
+        self.assertEqual(test.a, 'a')
+        self.assertEqual(test.b, 'b')
+        self.assertEqual(test.c, 'c')
+
+        test = Test()
+        test.method(a=1, b=2)
+        self.assertEqual(test.a, 1)
+        self.assertEqual(test.b, 2)
+
+
+    def test_typo_enter(self):
+        class mycontext(ContextDecorator):
+            def __unter__(self):
+                pass
+            def __exit__(self, *exc):
+                pass
+
+        with self.assertRaises(AttributeError):
+            with mycontext():
+                pass
+
+
+    def test_typo_exit(self):
+        class mycontext(ContextDecorator):
+            def __enter__(self):
+                pass
+            def __uxit__(self, *exc):
+                pass
+
+        with self.assertRaises(AttributeError):
+            with mycontext():
+                pass
+
+
+    def test_contextdecorator_as_mixin(self):
+        class somecontext(object):
+            started = False
+            exc = None
+
+            def __enter__(self):
+                self.started = True
+                return self
+
+            def __exit__(self, *exc):
+                self.exc = exc
+
+        class mycontext(somecontext, ContextDecorator):
+            pass
+
+        context = mycontext()
+        @context
+        def test():
+            self.assertIsNone(context.exc)
+            self.assertTrue(context.started)
+        test()
+        self.assertEqual(context.exc, (None, None, None))
+
+
+    def test_contextmanager_as_decorator(self):
+        state = []
+        @contextmanager
+        def woohoo(y):
+            state.append(y)
+            yield
+            state.append(999)
+
+        @woohoo(1)
+        def test(x):
+            self.assertEqual(state, [1])
+            state.append(x)
+        test('something')
+        self.assertEqual(state, [1, 'something', 999])
+
+
 # This is needed to make the test actually run under regrtest.py!
 def test_main():
     support.run_unittest(__name__)
diff --git a/Misc/NEWS b/Misc/NEWS
index 4a91b13..5d3c14e 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -460,6 +460,10 @@ C-API
 Library
 -------
 
+- Issue #9110: Addition of ContextDecorator to contextlib, for creating APIs
+  that act as both context managers and decorators. contextmanager changes
+  to use ContextDecorator.
+
 - Implement importlib.abc.SourceLoader and deprecate PyLoader and PyPycLoader
   for removal in Python 3.4.