Added support for pickling BaseQueryset instances Added BaseQueryset.__getstate__, BaseQuerySet.__setstate__ methods
This commit is contained in:
		| @@ -123,6 +123,39 @@ class BaseQuerySet(object): | ||||
|  | ||||
|         return queryset | ||||
|  | ||||
|     def __getstate__(self): | ||||
|         """ | ||||
|         Need for pickling queryset | ||||
|  | ||||
|         See https://github.com/MongoEngine/mongoengine/issues/442 | ||||
|         """ | ||||
|  | ||||
|         obj_dict = self.__dict__.copy() | ||||
|  | ||||
|         # don't picke collection, instead pickle collection params | ||||
|         obj_dict.pop("_collection_obj") | ||||
|  | ||||
|         # don't pickle cursor | ||||
|         obj_dict["_cursor_obj"] = None | ||||
|  | ||||
|         return obj_dict | ||||
|  | ||||
|     def __setstate__(self, obj_dict): | ||||
|         """ | ||||
|         Need for pickling queryset | ||||
|  | ||||
|         See https://github.com/MongoEngine/mongoengine/issues/442 | ||||
|         """ | ||||
|  | ||||
|         obj_dict["_collection_obj"] = obj_dict["_document"]._get_collection() | ||||
|  | ||||
|         # update attributes | ||||
|         self.__dict__.update(obj_dict) | ||||
|  | ||||
|         # forse load cursor | ||||
|         self._cursor | ||||
|  | ||||
|  | ||||
|     def __getitem__(self, key): | ||||
|         """Support skip and limit using getitem and slicing syntax. | ||||
|         """ | ||||
|   | ||||
							
								
								
									
										78
									
								
								tests/queryset/pickable.py
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										78
									
								
								tests/queryset/pickable.py
									
									
									
									
									
										Normal file
									
								
							| @@ -0,0 +1,78 @@ | ||||
| import pickle | ||||
| import unittest | ||||
| from pymongo.mongo_client import MongoClient | ||||
| from mongoengine import Document, StringField, IntField | ||||
| from mongoengine.connection import connect | ||||
|  | ||||
| __author__ = 'stas' | ||||
|  | ||||
| class Person(Document): | ||||
|     name = StringField() | ||||
|     age = IntField() | ||||
|  | ||||
| class TestQuerysetPickable(unittest.TestCase): | ||||
|     """ | ||||
|     Test for adding pickling support for QuerySet instances | ||||
|     See issue https://github.com/MongoEngine/mongoengine/issues/442 | ||||
|     """ | ||||
|     def setUp(self): | ||||
|         super(TestQuerysetPickable, self).setUp() | ||||
|  | ||||
|         connection = connect(db="test") #type: pymongo.mongo_client.MongoClient | ||||
|  | ||||
|         connection.drop_database("test") | ||||
|  | ||||
|         self.john = Person.objects.create( | ||||
|             name="John", | ||||
|             age=21 | ||||
|         ) | ||||
|  | ||||
|  | ||||
|     def test_picke_simple_qs(self): | ||||
|  | ||||
|         qs = Person.objects.all() | ||||
|  | ||||
|         pickle.dumps(qs) | ||||
|  | ||||
|     def _get_loaded(self, qs): | ||||
|         s = pickle.dumps(qs) | ||||
|  | ||||
|         return pickle.loads(s) | ||||
|  | ||||
|     def test_unpickle(self): | ||||
|         qs = Person.objects.all() | ||||
|  | ||||
|         loadedQs = self._get_loaded(qs) | ||||
|  | ||||
|         self.assertEqual(qs.count(), loadedQs.count()) | ||||
|  | ||||
|         #can update loadedQs | ||||
|         loadedQs.update(age=23) | ||||
|  | ||||
|         #check | ||||
|         self.assertEqual(Person.objects.first().age, 23) | ||||
|  | ||||
|     def test_pickle_support_filtration(self): | ||||
|         Person.objects.create( | ||||
|             name="Alice", | ||||
|             age=22 | ||||
|         ) | ||||
|  | ||||
|         Person.objects.create( | ||||
|             name="Bob", | ||||
|             age=23 | ||||
|         ) | ||||
|  | ||||
|         qs = Person.objects.filter(age__gte=22) | ||||
|         self.assertEqual(qs.count(), 2) | ||||
|  | ||||
|         loaded = self._get_loaded(qs) | ||||
|  | ||||
|         self.assertEqual(loaded.count(), 2) | ||||
|         self.assertEqual(loaded.filter(name="Bob").first().age, 23) | ||||
|  | ||||
|  | ||||
|  | ||||
|  | ||||
|  | ||||
|  | ||||
		Reference in New Issue
	
	Block a user