Ticket #16326: patch-16326-niran-2.diff

django/template/response.py

@@ -4 +4 @@
 class ContentNotRenderedError(Exception):
     pass
 
+class DiscardedAttributeError(AttributeError):
+    pass
+
 class SimpleTemplateResponse(HttpResponse):
 
+    RENDERING_ATTRS = ['template_name', 'context_data',
+        '_post_render_callbacks']
+
     def __init__(self, template, context=None, mimetype=None, status=None,
             content_type=None):
         # It would seem obvious to call these next two members 'template' and
@@ -37 +43 @@
         obj_dict = self.__dict__.copy()
         if not self._is_rendered:
             raise ContentNotRenderedError('The response content must be rendered before it can be pickled.')
-        del obj_dict['template_name']
-        del obj_dict['context_data']
-        del obj_dict['_post_render_callbacks']
+        for attr in self.RENDERING_ATTRS:
+            if attr in obj_dict:
+                del obj_dict[attr]
 
         return obj_dict
 
+    def __getattr__(self, name):
+        if name in self.RENDERING_ATTRS:
+            raise DiscardedAttributeError('The %s attribute was discarded when this TemplateResponse was pickled.' % name)
+        return super(SimpleTemplateResponse, self).__getattr__(name)
+
     def resolve_template(self, template):
         "Accepts a template object, path-to-template or list of paths"
         if isinstance(template, (list, tuple)):
@@ -120 +131 @@
 
     content = property(_get_content, _set_content)
 
-
 class TemplateResponse(SimpleTemplateResponse):
+
+    RENDERING_ATTRS = SimpleTemplateResponse.RENDERING_ATTRS + \
+        ['_request', '_current_app']
+
     def __init__(self, request, template, context=None, mimetype=None,
             status=None, content_type=None, current_app=None):
         # self.request gets over-written by django.test.client.Client - and
@@ -134 +148 @@
         super(TemplateResponse, self).__init__(
             template, context, mimetype, status, content_type)
 
-    def __getstate__(self):
-        """Pickling support function.
-
-        Ensures that the object can't be pickled before it has been
-        rendered, and that the pickled state only includes rendered
-        data, not the data used to construct the response.
-        """
-        obj_dict = super(TemplateResponse, self).__getstate__()
-
-        del obj_dict['_request']
-        del obj_dict['_current_app']
-
-        return obj_dict
-
     def resolve_context(self, context):
         """Convert context data into a full RequestContext object
         (assuming it isn't already a Context object).

tests/regressiontests/templates/response.py

@@ -8 +8 @@
 import django.template.context
 from django.template import Template, Context, RequestContext
 from django.template.response import (TemplateResponse, SimpleTemplateResponse,
-                                      ContentNotRenderedError)
+                                      ContentNotRenderedError,
+                                      DiscardedAttributeError)
 
 def test_processor(request):
     return {'processors': 'yes'}
@@ -190 +191 @@
 
         # ...and the unpickled reponse doesn't have the
         # template-related attributes, so it can't be re-rendered
-        self.assertFalse(hasattr(unpickled_response, 'template_name'))
-        self.assertFalse(hasattr(unpickled_response, 'context_data'))
-        self.assertFalse(hasattr(unpickled_response, '_post_render_callbacks'))
+        template_attrs = ('template_name', 'context_data', '_post_render_callbacks')
+        for attr in template_attrs:
+            self.assertFalse(hasattr(unpickled_response, attr))
+
+        # ...and requesting any of those attributes raises an exception
+        for attr in template_attrs:
+            with self.assertRaises(DiscardedAttributeError) as cm:
+                getattr(unpickled_response, attr)
+
+    def test_repickling(self):
+        response = SimpleTemplateResponse('first/test.html', {
+                'value': 123,
+                'fn': datetime.now,
+            })
+        self.assertRaises(ContentNotRenderedError,
+                          pickle.dumps, response)
+
+        response.render()
+        pickled_response = pickle.dumps(response)
+        unpickled_response = pickle.loads(pickled_response)
+        repickled_response = pickle.dumps(unpickled_response)
 
 class TemplateResponseTest(BaseTemplateResponseTest):
 
@@ -255 +274 @@
 
         # ...and the unpickled reponse doesn't have the
         # template-related attributes, so it can't be re-rendered
-        self.assertFalse(hasattr(unpickled_response, '_request'))
-        self.assertFalse(hasattr(unpickled_response, 'template_name'))
-        self.assertFalse(hasattr(unpickled_response, 'context_data'))
-        self.assertFalse(hasattr(unpickled_response, '_post_render_callbacks'))
+        template_attrs = ('template_name', 'context_data',
+            '_post_render_callbacks', '_request', '_current_app')
+        for attr in template_attrs:
+            self.assertFalse(hasattr(unpickled_response, attr))
+
+        # ...and requesting any of those attributes raises an exception
+        for attr in template_attrs:
+            with self.assertRaises(DiscardedAttributeError) as cm:
+                getattr(unpickled_response, attr)
+
+    def test_repickling(self):
+        response = SimpleTemplateResponse('first/test.html', {
+                'value': 123,
+                'fn': datetime.now,
+            })
+        self.assertRaises(ContentNotRenderedError,
+                          pickle.dumps, response)
+
+        response.render()
+        pickled_response = pickle.dumps(response)
+        unpickled_response = pickle.loads(pickled_response)
+        repickled_response = pickle.dumps(unpickled_response)
 
 
 class CustomURLConfTest(TestCase):