From b9407b26dff47232d57133d93593a1a50eee1cc5 Mon Sep 17 00:00:00 2001
From: Malcolm Tredinnick <malcolm.tredinnick@gmail.com>
Date: Wed, 20 Aug 2008 22:38:15 +0000
Subject: [PATCH] Made it possible to pickle DateQuerySets.

git-svn-id: http://code.djangoproject.com/svn/django/trunk@8455 bcc190cf-cafb-0310-a4f2-bffc1f526a37
---
 django/db/models/query.py               | 19 +++++++++----------
 django/db/models/sql/datastructures.py  |  2 +-
 django/db/models/sql/subqueries.py      | 17 +++++++++++++++++
 tests/regressiontests/queries/models.py |  4 ++++
 4 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/django/db/models/query.py b/django/db/models/query.py
index 1bc84c47ca..d7180bbc1f 100644
--- a/django/db/models/query.py
+++ b/django/db/models/query.py
@@ -452,12 +452,8 @@ class QuerySet(object):
                 "'kind' must be one of 'year', 'month' or 'day'."
         assert order in ('ASC', 'DESC'), \
                 "'order' must be either 'ASC' or 'DESC'."
-        # Let the FieldDoesNotExist exception propagate.
-        field = self.model._meta.get_field(field_name, many_to_many=False)
-        assert isinstance(field, DateField), "%r isn't a DateField." \
-                % field_name
-        return self._clone(klass=DateQuerySet, setup=True, _field=field,
-                _kind=kind, _order=order)
+        return self._clone(klass=DateQuerySet, setup=True,
+                _field_name=field_name, _kind=kind, _order=order)
 
     def none(self):
         """
@@ -721,13 +717,16 @@ class DateQuerySet(QuerySet):
         """
         self.query = self.query.clone(klass=sql.DateQuery, setup=True)
         self.query.select = []
-        self.query.add_date_select(self._field, self._kind, self._order)
-        if self._field.null:
-            self.query.add_filter(('%s__isnull' % self._field.name, False))
+        field = self.model._meta.get_field(self._field_name, many_to_many=False)
+        assert isinstance(field, DateField), "%r isn't a DateField." \
+                % field_name
+        self.query.add_date_select(field, self._kind, self._order)
+        if field.null:
+            self.query.add_filter(('%s__isnull' % field.name, False))
 
     def _clone(self, klass=None, setup=False, **kwargs):
         c = super(DateQuerySet, self)._clone(klass, False, **kwargs)
-        c._field = self._field
+        c._field_name = self._field_name
         c._kind = self._kind
         if setup and hasattr(c, '_setup_query'):
             c._setup_query()
diff --git a/django/db/models/sql/datastructures.py b/django/db/models/sql/datastructures.py
index cb54a564f8..913d8fde25 100644
--- a/django/db/models/sql/datastructures.py
+++ b/django/db/models/sql/datastructures.py
@@ -85,7 +85,7 @@ class Date(object):
     def __init__(self, col, lookup_type, date_sql_func):
         self.col = col
         self.lookup_type = lookup_type
-        self.date_sql_func= date_sql_func
+        self.date_sql_func = date_sql_func
 
     def relabel_aliases(self, change_map):
         c = self.col
diff --git a/django/db/models/sql/subqueries.py b/django/db/models/sql/subqueries.py
index 5ca041cbde..cea2a58d39 100644
--- a/django/db/models/sql/subqueries.py
+++ b/django/db/models/sql/subqueries.py
@@ -343,6 +343,23 @@ class DateQuery(Query):
     date field. This requires some special handling when converting the results
     back to Python objects, so we put it in a separate class.
     """
+    def __getstate__(self):
+        """
+        Special DateQuery-specific pickle handling.
+        """
+        for elt in self.select:
+            if isinstance(elt, Date):
+                # Eliminate a method reference that can't be pickled. The
+                # __setstate__ method restores this.
+                elt.date_sql_func = None
+        return super(DateQuery, self).__getstate__()
+
+    def __setstate__(self, obj_dict):
+        super(DateQuery, self).__setstate__(obj_dict)
+        for elt in self.select:
+            if isinstance(elt, Date):
+                self.date_sql_func = self.connection.ops.date_trunc_sql
+
     def results_iter(self):
         """
         Returns an iterator over the results from executing this query.
diff --git a/tests/regressiontests/queries/models.py b/tests/regressiontests/queries/models.py
index 60a3cfe2ca..debce6ad50 100644
--- a/tests/regressiontests/queries/models.py
+++ b/tests/regressiontests/queries/models.py
@@ -897,6 +897,10 @@ Bug #8283 -- Checking that applying filters after a disjunction works correctly.
 >>> (ExtraInfo.objects.filter(info='e2')|ExtraInfo.objects.filter(note=n1)).filter(note=n1)
 [<ExtraInfo: e1>]
 
+Pickling of DateQuerySets used to fail
+>>> qs = Item.objects.dates('created', 'month')
+>>> _ = pickle.loads(pickle.dumps(qs))
+
 """}
 
 # In Python 2.3, exceptions raised in __len__ are swallowed (Python issue