Allowed Context.push to behave as a context mananger.

Thanks Loic Bistuer for the review.

django/template/context.py

@@ -12,6 +12,21 @@ class ContextPopException(Exception):
     "pop() has been called more times than push()"
     pass
 
+
+class ContextDict(dict):
+    def __init__(self, context, *args, **kwargs):
+        super(ContextDict, self).__init__(*args, **kwargs)
+
+        context.dicts.append(self)
+        self.context = context
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        self.context.pop()
+
+
 class BaseContext(object):
     def __init__(self, dict_=None):
         self._reset_dicts(dict_)
@@ -34,10 +49,8 @@ class BaseContext(object):
         for d in reversed(self.dicts):
             yield d
 
-    def push(self):
-        d = {}
-        self.dicts.append(d)
-        return d
+    def push(self, *args, **kwargs):
+        return ContextDict(self, *args, **kwargs)
 
     def pop(self):
         if len(self.dicts) == 1:
@@ -83,6 +96,7 @@ class BaseContext(object):
         new_context._reset_dicts(values)
         return new_context
 
+
 class Context(BaseContext):
     "A stack container for variable context"
     def __init__(self, dict_=None, autoescape=True, current_app=None,
@@ -106,6 +120,7 @@ class Context(BaseContext):
         self.dicts.append(other_dict)
         return other_dict
 
+
 class RenderContext(BaseContext):
     """
     A stack container for storing Template state.

django/template/defaulttags.py

@@ -95,10 +95,9 @@ class FilterNode(Node):
     def render(self, context):
         output = self.nodelist.render(context)
         # Apply filters.
-        context.update({'var': output})
-        filtered = self.filter_expr.resolve(context)
-        context.pop()
-        return filtered
+        with context.push(var=output):
+            return self.filter_expr.resolve(context)
+
 
 class FirstOfNode(Node):
     def __init__(self, variables, escape=False):
@@ -143,7 +142,7 @@ class ForNode(Node):
             parentloop = context['forloop']
         else:
             parentloop = {}
-        context.push()
+        with context.push():
             try:
                 values = self.sequence.resolve(context, True)
             except VariableDoesNotExist:
@@ -154,7 +153,6 @@ class ForNode(Node):
                 values = list(values)
             len_values = len(values)
             if len_values < 1:
-            context.pop()
                 return self.nodelist_empty.render(context)
             nodelist = NodeList()
             if self.is_reversed:
@@ -207,7 +205,6 @@ class ForNode(Node):
                     # don't want to leave any vars from the previous loop on the
                     # context.
                     context.pop()
-        context.pop()
         return nodelist.render(context)
 
 class IfChangedNode(Node):
@@ -500,10 +497,9 @@ class WithNode(Node):
     def render(self, context):
         values = dict([(key, val.resolve(context)) for key, val in
                        six.iteritems(self.extra_context)])
-        context.update(values)
-        output = self.nodelist.render(context)
-        context.pop()
-        return output
+        with context.push(**values):
+            return self.nodelist.render(context)
+
 
 @register.tag
 def autoescape(parser, token):

django/template/loader.py

@@ -164,11 +164,8 @@ def render_to_string(template_name, dictionary=None, context_instance=None):
         return t.render(Context(dictionary))
     # Add the dictionary to the context stack, ensuring it gets removed again
     # to keep the context_instance in the same state it started in.
-    context_instance.update(dictionary)
-    try:
+    with context_instance.push(dictionary):
         return t.render(context_instance)
-    finally:
-        context_instance.pop()
 
 def select_template(template_name_list):
     "Given a list of template names, returns the first that can be loaded."

django/template/loader_tags.py

@@ -47,7 +47,7 @@ class BlockNode(Node):
 
     def render(self, context):
         block_context = context.render_context.get(BLOCK_CONTEXT_KEY)
-        context.push()
+        with context.push():
             if block_context is None:
                 context['block'] = self
                 result = self.nodelist.render(context)
@@ -62,7 +62,6 @@ class BlockNode(Node):
                 result = block.nodelist.render(context)
                 if push is not None:
                     block_context.push(self.name, push)
-        context.pop()
         return result
 
     def super(self):
@@ -133,10 +132,9 @@ class BaseIncludeNode(Node):
                        in six.iteritems(self.extra_context)])
         if self.isolated_context:
             return template.render(context.new(values))
-        context.update(values)
-        output = template.render(context)
-        context.pop()
-        return output
+        with context.push(**values):
+            return template.render(context)
+
 
 class ConstantIncludeNode(BaseIncludeNode):
     def __init__(self, template_path, *args, **kwargs):

docs/ref/templates/api.txt

@@ -325,6 +325,31 @@ If you ``pop()`` too much, it'll raise
     ...
     django.template.ContextPopException
 
+.. versionadded:: 1.7
+
+You can also use ``push()`` as a context manager to ensure a matching ``pop()``
+is called.
+
+    >>> c = Context()
+    >>> c['foo'] = 'first level'
+    >>> with c.push():
+    >>>     c['foo'] = 'second level'
+    >>>     c['foo']
+    'second level'
+    >>> c['foo']
+    'first level'
+
+All arguments passed to ``push()`` will be passed to the ``dict`` constructor
+used to build the new context level.
+
+    >>> c = Context()
+    >>> c['foo'] = 'first level'
+    >>> with c.push(foo='second level'):
+    >>>     c['foo']
+    'second level'
+    >>> c['foo']
+    'first level'
+
 .. method:: update(other_dict)
 
 In addition to ``push()`` and ``pop()``, the ``Context``

docs/releases/1.7.txt

@@ -60,6 +60,13 @@ Minor features
 * :attr:`~django.db.models.Options.app_label` is no longer required for models
   that are defined in a ``models`` package within an app.
 
+* The :meth:`Context.push() <django.template.Context.push>` method now returns
+  a context manager which automatically calls :meth:`pop()
+  <django.template.Context.pop>` upon exiting the ``with`` statement.
+  Additionally, :meth:`push() <django.template.Context.push>` now accepts
+  parameters that are passed to the ``dict`` constructor used to build the new
+  context level.
+
 Backwards incompatible changes in 1.7
 =====================================
 

tests/template_tests/test_context.py

@@ -16,3 +16,12 @@ class ContextTests(TestCase):
         self.assertEqual(c.pop(), {"a": 2})
         self.assertEqual(c["a"], 1)
         self.assertEqual(c.get("foo", 42), 42)
+
+        with c.push():
+            c['a'] = 2
+            self.assertEqual(c['a'], 2)
+        self.assertEqual(c['a'], 1)
+
+        with c.push(a=3):
+            self.assertEqual(c['a'], 3)
+        self.assertEqual(c['a'], 1)