Add GenericLazyReferenceField

docs/changelog.rst

@@ -4,7 +4,7 @@ Changelog
 
 Changes in 0.15.0
 =================
-- Add LazyReferenceField to address #1230
+- Add LazyReferenceField and GenericLazyReferenceField to address #1230
 
 Changes in 0.14.1
 =================

mongoengine/fields.py

@@ -47,8 +47,7 @@ __all__ = (
     'GenericEmbeddedDocumentField', 'DynamicField', 'ListField',
     'SortedListField', 'EmbeddedDocumentListField', 'DictField',
     'MapField', 'ReferenceField', 'CachedReferenceField',
-    'LazyReferenceField',
+    'LazyReferenceField', 'GenericLazyReferenceField',
-    # 'GenericLazyReferenceField',
     'GenericReferenceField', 'BinaryField', 'GridFSError', 'GridFSProxy',
     'FileField', 'ImageGridFsProxy', 'ImproperlyConfigured', 'ImageField',
     'GeoPointField', 'PointField', 'LineStringField', 'PolygonField',
@@ -1275,6 +1274,12 @@ class GenericReferenceField(BaseField):
     """A reference to *any* :class:`~mongoengine.document.Document` subclass
     that will be automatically dereferenced on access (lazily).
 
+    Note this field works the same way as :class:`~mongoengine.document.ReferenceField`,
+    doing database I/O access the first time it is accessed (even if it's to access
+    it ``pk`` or ``id`` field).
+    To solve this you should consider using the
+    :class:`~mongoengine.fields.GenericLazyReferenceField`.
+
     .. note ::
         * Any documents used as a generic reference must be registered in the
           document registry.  Importing the model will automatically register
@@ -2159,6 +2164,8 @@ class LazyReferenceField(BaseField):
     """A really lazy reference to a document.
     Unlike the :class:`~mongoengine.fields.ReferenceField` it must be manually
     dereferenced using it ``fetch()`` method.
+
+    .. versionadded:: 0.15
     """
 
     def __init__(self, document_type, passthrough=False, dbref=False,
@@ -2274,3 +2281,65 @@ class LazyReferenceField(BaseField):
 
     def lookup_member(self, member_name):
         return self.document_type._fields.get(member_name)
+
+
+class GenericLazyReferenceField(GenericReferenceField):
+    """A reference to *any* :class:`~mongoengine.document.Document` subclass
+    that will be automatically dereferenced on access (lazily).
+    Unlike the :class:`~mongoengine.fields.GenericReferenceField` it must be
+    manually dereferenced using it ``fetch()`` method.
+
+    .. note ::
+        * Any documents used as a generic reference must be registered in the
+          document registry.  Importing the model will automatically register
+          it.
+
+        * You can use the choices param to limit the acceptable Document types
+
+    .. versionadded:: 0.15
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.passthrough = kwargs.pop('passthrough', False)
+        super(GenericLazyReferenceField, self).__init__(*args, **kwargs)
+
+    def _validate_choices(self, value):
+        if isinstance(value, LazyReference):
+            value = value.document_type
+        super(GenericLazyReferenceField, self)._validate_choices(value)
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            return self
+
+        value = instance._data.get(self.name)
+        if isinstance(value, LazyReference):
+            if value.passthrough != self.passthrough:
+                instance._data[self.name] = LazyReference(
+                    value.document_type, value.pk, passthrough=self.passthrough)
+        elif value is not None:
+            if isinstance(value, (dict, SON)):
+                value = LazyReference(get_document(value['_cls']), value['_ref'].id, passthrough=self.passthrough)
+            elif isinstance(value, Document):
+                value = LazyReference(type(value), value.pk, passthrough=self.passthrough)
+            instance._data[self.name] = value
+
+        return super(GenericLazyReferenceField, self).__get__(instance, owner)
+
+    def validate(self, value):
+        if isinstance(value, LazyReference) and value.pk is None:
+            self.error('You can only reference documents once they have been'
+                       ' saved to the database')
+        return super(GenericLazyReferenceField, self).validate(value)
+
+    def to_mongo(self, document):
+        if document is None:
+            return None
+
+        if isinstance(document, LazyReference):
+            return SON((
+                ('_cls', document.document_type._class_name),
+                ('_ref', document)
+            ))
+        else:
+            return super(GenericLazyReferenceField, self).to_mongo(document)

tests/fields/fields.py

@@ -933,7 +933,7 @@ class FieldTest(MongoDBTestCase):
             authors = ListField(ReferenceField(User))
             authors_as_lazy = ListField(LazyReferenceField(User))
             generic = ListField(GenericReferenceField())
-            # generic_as_lazy = ListField(LazyGenericReferenceField())
+            generic_as_lazy = ListField(GenericLazyReferenceField())
 
         User.drop_collection()
         BlogPost.drop_collection()
@@ -992,17 +992,17 @@ class FieldTest(MongoDBTestCase):
         post.generic = [user]
         post.validate()
 
-        # post.generic_as_lazy = [1, 2]
+        post.generic_as_lazy = [1, 2]
-        # self.assertRaises(ValidationError, post.validate)
+        self.assertRaises(ValidationError, post.validate)
 
-        # post.generic_as_lazy = [User(), Comment()]
+        post.generic_as_lazy = [User(), Comment()]
-        # self.assertRaises(ValidationError, post.validate)
+        self.assertRaises(ValidationError, post.validate)
 
-        # post.generic_as_lazy = [Comment()]
+        post.generic_as_lazy = [Comment()]
-        # self.assertRaises(ValidationError, post.validate)
+        self.assertRaises(ValidationError, post.validate)
 
-        # post.generic_as_lazy = [user]
+        post.generic_as_lazy = [user]
-        # post.validate()
+        post.validate()
 
     def test_sorted_list_sorting(self):
         """Ensure that a sorted list field properly sorts values.
@@ -4872,5 +4872,185 @@ class LazyReferenceFieldTest(MongoDBTestCase):
         self.assertNotEqual(other_animalref, animal)
 
 
+class GenericLazyReferenceFieldTest(MongoDBTestCase):
+    def test_generic_lazy_reference_simple(self):
+        class Animal(Document):
+            name = StringField()
+            tag = StringField()
+
+        class Ocurrence(Document):
+            person = StringField()
+            animal = GenericLazyReferenceField()
+
+        Animal.drop_collection()
+        Ocurrence.drop_collection()
+
+        animal = Animal(name="Leopard", tag="heavy").save()
+        Ocurrence(person="test", animal=animal).save()
+        p = Ocurrence.objects.get()
+        self.assertIsInstance(p.animal, LazyReference)
+        fetched_animal = p.animal.fetch()
+        self.assertEqual(fetched_animal, animal)
+        # `fetch` keep cache on referenced document by default...
+        animal.tag = "not so heavy"
+        animal.save()
+        double_fetch = p.animal.fetch()
+        self.assertIs(fetched_animal, double_fetch)
+        self.assertEqual(double_fetch.tag, "heavy")
+        # ...unless specified otherwise
+        fetch_force = p.animal.fetch(force=True)
+        self.assertIsNot(fetch_force, fetched_animal)
+        self.assertEqual(fetch_force.tag, "not so heavy")
+
+    def test_generic_lazy_reference_choices(self):
+        class Animal(Document):
+            name = StringField()
+
+        class Vegetal(Document):
+            name = StringField()
+
+        class Mineral(Document):
+            name = StringField()
+
+        class Ocurrence(Document):
+            living_thing = GenericLazyReferenceField(choices=[Animal, Vegetal])
+            thing = GenericLazyReferenceField()
+
+        Animal.drop_collection()
+        Vegetal.drop_collection()
+        Mineral.drop_collection()
+        Ocurrence.drop_collection()
+
+        animal = Animal(name="Leopard").save()
+        vegetal = Vegetal(name="Oak").save()
+        mineral = Mineral(name="Granite").save()
+
+        occ_animal = Ocurrence(living_thing=animal, thing=animal).save()
+        occ_vegetal = Ocurrence(living_thing=vegetal, thing=vegetal).save()
+        with self.assertRaises(ValidationError):
+            Ocurrence(living_thing=mineral).save()
+
+        occ = Ocurrence.objects.get(living_thing=animal)
+        self.assertEqual(occ, occ_animal)
+        self.assertIsInstance(occ.thing, LazyReference)
+        self.assertIsInstance(occ.living_thing, LazyReference)
+
+        occ.thing = vegetal
+        occ.living_thing = vegetal
+        occ.save()
+
+        occ.thing = mineral
+        occ.living_thing = mineral
+        with self.assertRaises(ValidationError):
+            occ.save()
+
+    def test_generic_lazy_reference_set(self):
+        class Animal(Document):
+            meta = {'allow_inheritance': True}
+
+            name = StringField()
+            tag = StringField()
+
+        class Ocurrence(Document):
+            person = StringField()
+            animal = GenericLazyReferenceField()
+
+        Animal.drop_collection()
+        Ocurrence.drop_collection()
+
+        class SubAnimal(Animal):
+            nick = StringField()
+
+        animal = Animal(name="Leopard", tag="heavy").save()
+        sub_animal = SubAnimal(nick='doggo', name='dog').save()
+        for ref in (
+                animal,
+                LazyReference(Animal, animal.pk),
+                {'_cls': 'Animal', '_ref': DBRef(animal._get_collection_name(), animal.pk)},
+
+                sub_animal,
+                LazyReference(SubAnimal, sub_animal.pk),
+                {'_cls': 'SubAnimal', '_ref': DBRef(sub_animal._get_collection_name(), sub_animal.pk)},
+                ):
+            p = Ocurrence(person="test", animal=ref).save()
+            p.reload()
+            self.assertIsInstance(p.animal, (LazyReference, Document))
+            p.animal.fetch()
+
+    def test_generic_lazy_reference_bad_set(self):
+        class Animal(Document):
+            name = StringField()
+            tag = StringField()
+
+        class Ocurrence(Document):
+            person = StringField()
+            animal = GenericLazyReferenceField(choices=['Animal'])
+
+        Animal.drop_collection()
+        Ocurrence.drop_collection()
+
+        class BadDoc(Document):
+            pass
+
+        animal = Animal(name="Leopard", tag="heavy").save()
+        baddoc = BadDoc().save()
+        for bad in (
+                42,
+                'foo',
+                baddoc,
+                LazyReference(BadDoc, animal.pk)
+                ):
+            with self.assertRaises(ValidationError):
+                p = Ocurrence(person="test", animal=bad).save()
+
+    def test_generic_lazy_reference_query_conversion(self):
+        class Member(Document):
+            user_num = IntField(primary_key=True)
+
+        class BlogPost(Document):
+            title = StringField()
+            author = GenericLazyReferenceField()
+
+        Member.drop_collection()
+        BlogPost.drop_collection()
+
+        m1 = Member(user_num=1)
+        m1.save()
+        m2 = Member(user_num=2)
+        m2.save()
+
+        post1 = BlogPost(title='post 1', author=m1)
+        post1.save()
+
+        post2 = BlogPost(title='post 2', author=m2)
+        post2.save()
+
+        post = BlogPost.objects(author=m1).first()
+        self.assertEqual(post.id, post1.id)
+
+        post = BlogPost.objects(author=m2).first()
+        self.assertEqual(post.id, post2.id)
+
+        # Same thing by passing a LazyReference instance
+        post = BlogPost.objects(author=LazyReference(Member, m2.pk)).first()
+        self.assertEqual(post.id, post2.id)
+
+    def test_generic_lazy_reference_not_set(self):
+        class Animal(Document):
+            name = StringField()
+            tag = StringField()
+
+        class Ocurrence(Document):
+            person = StringField()
+            animal = GenericLazyReferenceField()
+
+        Animal.drop_collection()
+        Ocurrence.drop_collection()
+
+        Ocurrence(person='foo').save()
+        p = Ocurrence.objects.get()
+        self.assertIs(p.animal, None)
+
+
 if __name__ == '__main__':
     unittest.main()