Merge pull request #1520 from ZoetropeLabs/fix/allow-reference-fields-take-object-ids

Allow ReferenceFields to take ObjectIds

Author: Omer Katz

mongoengine/fields.py

@@ -998,8 +998,8 @@ def prepare_query_value(self, op, value):
 
     def validate(self, value):
 
-        if not isinstance(value, (self.document_type, DBRef)):
-            self.error('A ReferenceField only accepts DBRef or documents')
+        if not isinstance(value, (self.document_type, DBRef, ObjectId)):
+            self.error('A ReferenceField only accepts DBRef, ObjectId or documents')
 
         if isinstance(value, Document) and value.id is None:
             self.error('You can only reference documents once they have been '

tests/fields/fields.py

@@ -2088,6 +2088,12 @@ class BlogPost(Document):
         post1.author = post2
         self.assertRaises(ValidationError, post1.validate)
 
+        # Ensure ObjectID's are accepted as references
+        user_object_id = user.pk
+        post3 = BlogPost(content="Chips and curry sauce taste good.")
+        post3.author = user_object_id
+        post3.save()
+
         # Make sure referencing a saved document of the right type works
         user.save()
         post1.author = user
@@ -2098,6 +2104,20 @@ class BlogPost(Document):
         post1.author = post2
         self.assertRaises(ValidationError, post1.validate)
 
+    def test_objectid_reference_fields(self):
+        """Make sure storing Object ID references works."""
+        class Person(Document):
+            name = StringField()
+            parent = ReferenceField('self')
+
+        Person.drop_collection()
+
+        p1 = Person(name="John").save()
+        Person(name="Ross", parent=p1.pk).save()
+
+        p = Person.objects.get(name="Ross")
+        self.assertEqual(p.parent, p1)
+
     def test_dbref_reference_fields(self):
         """Make sure storing references as bson.dbref.DBRef works."""
         class Person(Document):