Added recursive / document name references

mongoengine/fields.py

@@ -13,6 +13,8 @@ __all__ = ['StringField', 'IntField', 'FloatField', 'BooleanField',
            'ObjectIdField', 'ReferenceField', 'ValidationError',
            'DecimalField', 'URLField', 'GenericReferenceField']
 
+RECURSIVE_REFERENCE_CONSTANT = 'self'
+
 
 class StringField(BaseField):
     """A unicode string field.
@@ -334,10 +336,19 @@ class ReferenceField(BaseField):
             if not issubclass(document_type, (Document, basestring)):
                 raise ValidationError('Argument to ReferenceField constructor '
                                       'must be a document class or a string')
-        self.document_type = document_type
+        self.document_type_obj = document_type
         self.document_obj = None
         super(ReferenceField, self).__init__(**kwargs)
 
+    @property
+    def document_type(self):
+        if isinstance(self.document_type_obj, basestring):
+            if self.document_type_obj == RECURSIVE_REFERENCE_CONSTANT:
+                self.document_type_obj = self.owner_document
+            else:
+                self.document_type_obj = get_document(self.document_type_obj)
+        return self.document_type_obj
+
     def __get__(self, instance, owner):
         """Descriptor to allow lazy dereferencing.
         """

tests/fields.py

@@ -350,6 +350,40 @@ class FieldTest(unittest.TestCase):
         User.drop_collection()
         Group.drop_collection()
 
+    def test_recursive_reference(self):
+        """Ensure that ReferenceFields can reference their own documents.
+        """
+        class Employee(Document):
+            name = StringField()
+            boss = ReferenceField('self')
+
+        bill = Employee(name='Bill Lumbergh')
+        bill.save()
+        peter = Employee(name='Peter Gibbons', boss=bill)
+        peter.save()
+
+        peter = Employee.objects.with_id(peter.id)
+        self.assertEqual(peter.boss, bill)
+
+    def test_undefined_reference(self):
+        """Ensure that ReferenceFields may reference undefined Documents.
+        """
+        class Product(Document):
+            name = StringField()
+            company = ReferenceField('Company')
+
+        class Company(Document):
+            name = StringField()
+
+        ten_gen = Company(name='10gen')
+        ten_gen.save()
+        mongodb = Product(name='MongoDB', company=ten_gen)
+        mongodb.save()
+
+        obj = Product.objects(company=ten_gen).first()
+        self.assertEqual(obj, mongodb)
+        self.assertEqual(obj.company, ten_gen)
+
     def test_reference_query_conversion(self):
         """Ensure that ReferenceFields can be queried using objects and values
         of the type of the primary key of the referenced object.