Added queryset_manager decorator

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