Queries now return correct subclasses of Documents

mongomap/base.py

@@ -97,10 +97,18 @@ class DocumentMetaclass(type):
             return super_new(cls, name, bases, attrs)
 
         doc_fields = {}
-        # Include all fields present in superclasses
+        class_name = [name]
+        superclasses = {}
         for base in bases:
+            # Include all fields present in superclasses
             if hasattr(base, '_fields'):
                 doc_fields.update(base._fields)
+                class_name.append(base._class_name)
+                # Get superclasses from superclass
+                superclasses[base._class_name] = base
+                superclasses.update(base._superclasses)
+        attrs['_class_name'] = '.'.join(reversed(class_name))
+        attrs['_superclasses'] = superclasses
 
         # Add the document's fields to the _fields attribute
         for attr_name, attr_value in attrs.items():
@@ -164,6 +172,21 @@ class BaseDocument(object):
                 # Use default value
                 setattr(self, attr_name, getattr(self, attr_name, None))
 
+    @classmethod
+    def _get_subclasses(cls):
+        """Return a dictionary of all subclasses (found recursively).
+        """
+        try:
+            subclasses = cls.__subclasses__()
+        except:
+            subclasses = cls.__subclasses__(cls)
+
+        all_subclasses = {}
+        for subclass in subclasses:
+            all_subclasses[subclass._class_name] = subclass
+            all_subclasses.update(subclass._get_subclasses())
+        return all_subclasses
+
     def __iter__(self):
         # Use _data rather than _fields as iterator only looks at names so
         # values don't need to be converted to Python types
@@ -203,12 +226,25 @@ class BaseDocument(object):
             value = getattr(self, field_name, None)
             if value is not None:
                 data[field_name] = field._to_mongo(value)
+        data['_cls'] = self._class_name
+        data['_types'] = self._superclasses.keys() + [self._class_name]
         return data
     
     @classmethod
     def _from_son(cls, son):
         """Create an instance of a Document (subclass) from a PyMongo SOM.
         """
+        class_name = son[u'_cls']
         data = dict((str(key), value) for key, value in son.items())
+        del data['_cls']
+
+        # Return correct subclass for document type
+        if class_name != cls._class_name:
+            subclasses = cls._get_subclasses()
+            if class_name not in subclasses:
+                # Type of document is probably more generic than the class
+                # that has been queried to return this SON
+                return None
+            cls = subclasses[class_name]
         return cls(**data)
 

mongomap/collection.py

@@ -76,13 +76,13 @@ class CollectionManager(object):
     def find(self, **query):
         """Query the collection for document matching the provided query.
         """
-        if query:
-            query = self._transform_query(**query)
+        query = self._transform_query(**query)
+        query['_types'] = self._document._class_name
         return QuerySet(self._document, self._collection.find(query))
 
     def find_one(self, **query):
         """Query the collection for document matching the provided query.
         """
-        if query:
-            query = self._transform_query(**query)
+        query = self._transform_query(**query)
+        query['_types'] = self._document._class_name
         return self._document._from_son(self._collection.find_one(query))

mongomap/fields.py

@@ -69,6 +69,7 @@ class EmbeddedDocumentField(BaseField):
     
     def _to_python(self, value):
         if not isinstance(value, self.document):
+            assert(isinstance(value, (dict, pymongo.son.SON)))
             return self.document._from_son(value)
         return value
 
@@ -107,11 +108,6 @@ class ListField(BaseField):
     def _validate(self, value):
         """Make sure that a list of valid fields is being used.
         """
-#        print
-#        print value
-#        print type(value)
-#        print isinstance(value, list)
-#        print
         if not isinstance(value, (list, tuple)):
             raise ValidationError('Only lists and tuples may be used in a '
                                   'list field')