Added queryset_manager decorator

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 <http://hmarr.com/mongoengine/tutorial.html>`_, a `user guide 
 <http://hmarr.com/mongoengine/userguide.html>`_ and an `API reference
 <http://hmarr.com/mongoengine/apireference.html>`_.
 
 Dependencies
 ============
-pymongo 1.1+
-sphinx (optional -- for documentation generation)
+- pymongo 1.1+
+- sphinx (optional - for documentation generation)

docs/apireference.rst

@@ -27,6 +27,8 @@ Querying
 .. autoclass:: mongoengine.queryset.QuerySet
    :members:
 
+.. autofunction:: mongoengine.queryset.queryset_manager
+
 Fields
 ======
 

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'
 

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
 

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

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 
@@ -197,5 +197,21 @@ class QuerySetManager(object):
             # Document class being used rather than a document object
             return self
 
-        # self._document should be the same as owner
-        return QuerySet(self._document, self._collection)
+        if self._collection is None:
+            db = _get_db()
+            self._collection = db[owner._meta['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)

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()