Factorize some code using ContextDecorator.

3 files changed, 24 insertions, 19 deletions

diff --git a/django/db/transaction.py b/django/db/transaction.py
index 4fb58d6dd8..abc7892975 100644
--- a/django/db/transaction.py
+++ b/django/db/transaction.py
@@ -1,9 +1,7 @@
-from functools import wraps
-
 from django.db import (
     connections, DEFAULT_DB_ALIAS,
     DatabaseError, Error, ProgrammingError)
-from django.utils.decorators import available_attrs
+from django.utils.decorators import ContextDecorator
 
 
 class TransactionManagementError(ProgrammingError):
@@ -109,7 +107,7 @@ def set_rollback(rollback, using=None):
 # Decorators / context managers #
 #################################
 
-class Atomic(object):
+class Atomic(ContextDecorator):
     """
     This class guarantees the atomic execution of a given block.
 
@@ -285,13 +283,6 @@ class Atomic(object):
                 else:
                     connection.in_atomic_block = False
 
-    def __call__(self, func):
-        @wraps(func, assigned=available_attrs(func))
-        def inner(*args, **kwargs):
-            with self:
-                return func(*args, **kwargs)
-        return inner
-
 
 def atomic(using=None, savepoint=True):
     # Bare decorator: @atomic -- although the first argument is called
diff --git a/django/test/utils.py b/django/test/utils.py
index 4c5c910be0..b2d9f7b696 100644
--- a/django/test/utils.py
+++ b/django/test/utils.py
@@ -18,6 +18,7 @@ from django.template import Template, loader, TemplateDoesNotExist
 from django.template.loaders import cached
 from django.test.signals import template_rendered, setting_changed
 from django.utils import six
+from django.utils.decorators import ContextDecorator
 from django.utils.deprecation import RemovedInDjango19Warning, RemovedInDjango20Warning
 from django.utils.encoding import force_str
 from django.utils.translation import deactivate
@@ -146,7 +147,7 @@ def get_runner(settings, test_runner_class=None):
     return test_runner
 
 
-class override_template_loaders(object):
+class override_template_loaders(ContextDecorator):
     """
     Acts as a function decorator, context manager or start/end manager and
     override the template loaders. It could be used in the following ways:
@@ -174,13 +175,6 @@ class override_template_loaders(object):
     def __exit__(self, type, value, traceback):
         loader.template_source_loaders = self.old_loaders
 
-    def __call__(self, test_func):
-        @wraps(test_func)
-        def inner(*args, **kwargs):
-            with self:
-                return test_func(*args, **kwargs)
-        return inner
-
     @classmethod
     def override(cls, *loaders):
         if hasattr(loader, RESTORE_LOADERS_ATTR):
diff --git a/django/utils/decorators.py b/django/utils/decorators.py
index b338c8026c..3331fca076 100644
--- a/django/utils/decorators.py
+++ b/django/utils/decorators.py
@@ -1,5 +1,10 @@
 "Functions that help with dynamically creating decorators for views."
 
+try:
+    from contextlib import ContextDecorator
+except ImportError:
+    ContextDecorator = None
+
 from functools import wraps, update_wrapper, WRAPPER_ASSIGNMENTS
 
 from django.utils import six
@@ -124,3 +129,18 @@ def make_middleware_decorator(middleware_class):
             return _wrapped_view
         return _decorator
     return _make_decorator
+
+
+if ContextDecorator is None:
+    # ContextDecorator was introduced in Python 3.2
+    # See https://docs.python.org/3/library/contextlib.html#contextlib.ContextDecorator
+    class ContextDecorator(object):
+        """
+        A base class that enables a context manager to also be used as a decorator.
+        """
+        def __call__(self, func):
+            @wraps(func, assigned=available_attrs(func))
+            def inner(*args, **kwargs):
+                with self:
+                    return func(*args, **kwargs)
+            return inner