diff --git a/README.rst b/README.rst index 28b2dc67..5b4b7094 100644 --- a/README.rst +++ b/README.rst @@ -7,12 +7,12 @@ MongoEngine About ===== MongoEngine is a Python Object-Document Mapper for working with MongoDB. -Documentation available at http://hmarr.com/mongoengine/ -- there is currently +Documentation available at http://hmarr.com/mongoengine/ - there is currently a `tutorial `_, a `user guide `_ and an `API reference `_. Dependencies ============ -pymongo 1.1+ -sphinx (optional -- for documentation generation) +- pymongo 1.1+ +- sphinx (optional - for documentation generation) diff --git a/docs/apireference.rst b/docs/apireference.rst index ea8411e2..86818805 100644 --- a/docs/apireference.rst +++ b/docs/apireference.rst @@ -27,6 +27,8 @@ Querying .. autoclass:: mongoengine.queryset.QuerySet :members: +.. autofunction:: mongoengine.queryset.queryset_manager + Fields ====== diff --git a/mongoengine/__init__.py b/mongoengine/__init__.py index ae6b7c19..ede2736f 100644 --- a/mongoengine/__init__.py +++ b/mongoengine/__init__.py @@ -4,8 +4,11 @@ import fields from fields import * import connection from connection import * +import queryset +from queryset import * -__all__ = document.__all__ + fields.__all__ + connection.__all__ +__all__ = (document.__all__ + fields.__all__ + connection.__all__ + + queryset.__all__) __author__ = 'Harry Marr' diff --git a/mongoengine/base.py b/mongoengine/base.py index 60409127..2e6fe52a 100644 --- a/mongoengine/base.py +++ b/mongoengine/base.py @@ -169,7 +169,7 @@ class TopLevelDocumentMetaclass(DocumentMetaclass): # Set up collection manager, needs the class to have fields so use # DocumentMetaclass before instantiating CollectionManager object new_class = super_new(cls, name, bases, attrs) - new_class.objects = QuerySetManager(new_class) + new_class.objects = QuerySetManager() return new_class diff --git a/mongoengine/document.py b/mongoengine/document.py index a4b78619..49ff238a 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -44,8 +44,8 @@ class Document(BaseDocument): document already exists, it will be updated, otherwise it will be created. """ - object_id = self.objects._collection.save(self.to_mongo()) - self.id = object_id + object_id = self.__class__.objects._collection.save(self.to_mongo()) + self.id = self._fields['id'].to_python(object_id) def delete(self): """Delete the :class:`~mongoengine.Document` from the database. This diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 68c27067..183f0eee 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -3,6 +3,9 @@ from connection import _get_db import pymongo +__all__ = ['queryset_manager'] + + class QuerySet(object): """A set of results returned from a query. Wraps a MongoDB cursor, providing :class:`~mongoengine.Document` objects as the results. @@ -182,12 +185,9 @@ class QuerySet(object): class QuerySetManager(object): - def __init__(self, document): - db = _get_db() - self._document = document - self._collection_name = document._meta['collection'] - # This will create the collection if it doesn't exist - self._collection = db[self._collection_name] + def __init__(self, manager_func=None): + self._manager_func = manager_func + self._collection = None def __get__(self, instance, owner): """Descriptor for instantiating a new QuerySet object when @@ -196,6 +196,22 @@ class QuerySetManager(object): if instance is not None: # Document class being used rather than a document object return self + + if self._collection is None: + db = _get_db() + self._collection = db[owner._meta['collection']] - # self._document should be the same as owner - return QuerySet(self._document, self._collection) + # owner is the document that contains the QuerySetManager + queryset = QuerySet(owner, self._collection) + if self._manager_func: + queryset = self._manager_func(queryset) + return queryset + +def queryset_manager(func): + """Decorator that allows you to define custom QuerySet managers on + :class:`~mongoengine.Document` classes. The manager must be a function that + accepts a :class:`~mongoengine.queryset.QuerySet` as its only argument, and + returns a :class:`~mongoengine.queryset.QuerySet`, probably the same one + but modified in some way. + """ + return QuerySetManager(func) diff --git a/tests/queryset.py b/tests/queryset.py index 461ad34b..34c363b8 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -214,6 +214,32 @@ class QuerySetTest(unittest.TestCase): BlogPost.drop_collection() + def test_custom_manager(self): + """Ensure that custom QuerySetManager instances work as expected. + """ + class BlogPost(Document): + tags = ListField(StringField()) + + @queryset_manager + def music_posts(queryset): + return queryset(tags='music') + + BlogPost.drop_collection() + + post1 = BlogPost(tags=['music', 'film']) + post1.save() + post2 = BlogPost(tags=['music']) + post2.save() + post3 = BlogPost(tags=['film', 'actors']) + post3.save() + + self.assertEqual([p.id for p in BlogPost.objects], + [post1.id, post2.id, post3.id]) + self.assertEqual([p.id for p in BlogPost.music_posts], + [post1.id, post2.id]) + + BlogPost.drop_collection() + def tearDown(self): self.Person.drop_collection()