Fixes SearchQuerySet not pickleable. Patch by oyiptong, tests by toas…

…tdriven.

AUTHORS

@@ -46,3 +46,4 @@ Thanks to
     * Rob Hudson (robhudson) for improvements to the admin search.
     * apollo13 for simplifying ``SearchForm.__init__``.
     * Carl Meyer (carljm) for a patch regarding character primary keys.
+    * oyiptong for a patch regarding pickling.

haystack/backends/__init__.py

@@ -297,9 +297,10 @@ def __getstate__(self):
         """For pickling."""
         obj_dict = self.__dict__.copy()
         del(obj_dict['backend'])
+
         # Rip off the class bits as we'll be using this path when we go to load
         # the backend.
-        obj_dict['backend_used'] = ".".join(str(self.backend).replace("<class '", "").replace("'>", "").split(".")[0:-1])
+        obj_dict['backend_used'] = ".".join(str(self.backend).replace("<", "").split(".")[0:-1])
         return obj_dict
 
     def __setstate__(self, obj_dict):

haystack/query.py

@@ -38,7 +38,16 @@ def __getstate__(self):
         len(self)
         obj_dict = self.__dict__.copy()
         obj_dict['_iter'] = None
+        del obj_dict['site']
         return obj_dict
+
+    def __setstate__(self, dict):
+        """
+        For unpickling.
+        """
+        self.__dict__ = dict
+        from haystack import site as main_site
+        self.site = main_site
 
     def __repr__(self):
         data = list(self[:REPR_OUTPUT_SIZE])

tests/core/tests/mocks.py

@@ -112,3 +112,8 @@ def run_mlt(self):
         results = self.backend.more_like_this(self._mlt_instance, final_query)
         self._results = results['results'][self.start_offset:self.end_offset]
         self._hit_count = results['hits']
+
+
+# For pickling tests.
+SearchBackend = MockSearchBackend
+SearchQuery = MockSearchQuery

tests/core/tests/query.py

@@ -18,6 +18,16 @@
 except NameError:
     from sets import Set as set
 
+test_pickling = True
+
+try:
+    import cPickle as pickle
+except ImportError:
+    try:
+        import pickle
+    except ImportError:
+        test_pickling = False
+
 
 class SQTestCase(TestCase):
     def test_split_expression(self):
@@ -661,3 +671,41 @@ def test_dictionary_lookup(self):
         EmptySearchQuerySets can be used in templates.
         """
         self.assertRaises(TypeError, lambda: self.esqs['count'])
+
+
+if test_pickling:
+    class PickleSearchQuerySetTestCase(TestCase):
+        def setUp(self):
+            super(PickleSearchQuerySetTestCase, self).setUp()
+            self.bsqs = SearchQuerySet(query=DummySearchQuery(backend=DummySearchBackend()))
+            self.msqs = SearchQuerySet(query=MockSearchQuery(backend=MockSearchBackend()))
+            self.mmsqs = SearchQuerySet(query=MockSearchQuery(backend=MixedMockSearchBackend()))
+
+            # Stow.
+            self.old_debug = settings.DEBUG
+            settings.DEBUG = True
+            self.old_site = haystack.site
+            test_site = SearchSite()
+            test_site.register(MockModel)
+            test_site.register(CharPKMockModel)
+            haystack.site = test_site
+
+            backends.reset_search_queries()
+
+        def tearDown(self):
+            # Restore.
+            haystack.site = self.old_site
+            settings.DEBUG = self.old_debug
+            super(PickleSearchQuerySetTestCase, self).tearDown()
+
+        def test_pickling(self):
+            results = self.msqs.all()
+
+            for res in results:
+                # Make sure the cache is full.
+                pass
+
+            in_a_pickle = pickle.dumps(results)
+            like_a_cuke = pickle.loads(in_a_pickle)
+            self.assertEqual(len(like_a_cuke), len(results))
+            self.assertEqual(like_a_cuke[0].id, results[0].id)

tests/solr_tests/tests/solr_backend.py

@@ -15,6 +15,16 @@
 except NameError:
     from sets import Set as set
 
+test_pickling = True
+
+try:
+    import cPickle as pickle
+except ImportError:
+    try:
+        import pickle
+    except ImportError:
+        test_pickling = False
+
 
 def clear_solr_index():
     # Wipe it clean.
@@ -794,3 +804,48 @@ def test_round_trip(self):
         self.assertEqual(result.created, datetime.datetime(2009, 11, 21, 21, 31, 00))
         self.assertEqual(result.tags, ['staff', 'outdoor', 'activist', 'scientist'])
         self.assertEqual(result.sites, [3, 5, 1])
+
+
+if test_pickling:
+    class LiveSolrPickleTestCase(TestCase):
+        fixtures = ['bulk_data.json']
+
+        def setUp(self):
+            super(LiveSolrPickleTestCase, self).setUp()
+
+            # Wipe it clean.
+            clear_solr_index()
+
+            # With the models registered, you get the proper bits.
+            import haystack
+            from haystack.sites import SearchSite
+
+            # Stow.
+            self.old_site = haystack.site
+            test_site = SearchSite()
+            test_site.register(MockModel, SolrMockModelSearchIndex)
+            test_site.register(AnotherMockModel, SolrAnotherMockModelSearchIndex)
+            haystack.site = test_site
+
+            self.sqs = SearchQuerySet()
+
+            test_site.get_index(MockModel).update()
+            test_site.get_index(AnotherMockModel).update()
+
+        def tearDown(self):
+            # Restore.
+            import haystack
+            haystack.site = self.old_site
+            super(LiveSolrPickleTestCase, self).tearDown()
+
+        def test_pickling(self):
+            results = self.sqs.all()
+
+            for res in results:
+                # Make sure the cache is full.
+                pass
+
+            in_a_pickle = pickle.dumps(results)
+            like_a_cuke = pickle.loads(in_a_pickle)
+            self.assertEqual(len(like_a_cuke), len(results))
+            self.assertEqual(like_a_cuke[0].id, results[0].id)