Merge pull request #303 from Scille/query_mapper

Allow passing Document and EmbeddedDocument in queries
Scille · Oct 12, 2020 · 7217b65 · 7217b65
2 parents 5b458d8 + 86a742c
commit 7217b65
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 1 deletion.
diff --git a/tests/test_query_mapper.py b/tests/test_query_mapper.py
@@ -1,5 +1,7 @@
 import datetime as dt
 
+from bson import ObjectId
+
 from umongo import Document, EmbeddedDocument, fields
 from umongo.query_mapper import map_query
 
@@ -10,6 +12,10 @@ class TestQueryMapper(BaseTest):
 
     def test_query_mapper(self):
 
+        @self.instance.register
+        class Editor(Document):
+            name = fields.StrField()
+
         @self.instance.register
         class Author(EmbeddedDocument):
             name = fields.StrField()
@@ -19,7 +25,6 @@ class Author(EmbeddedDocument):
         class Chapter(EmbeddedDocument):
             title = fields.StrField()
             pagination = fields.IntField(attribute='p')
-            # sub_chapters = fields.EmbeddedField('Chapter')
 
         @self.instance.register
         class Book(Document):
@@ -28,6 +33,7 @@ class Book(Document):
             author = fields.EmbeddedField(Author, attribute='a')
             chapters = fields.ListField(fields.EmbeddedField(Chapter))
             tags = fields.ListField(fields.StrField(), attribute='t')
+            editor = fields.ReferenceField(Editor, attribute='e')
 
         book_fields = Book.schema.fields
         # No changes needed
@@ -100,6 +106,22 @@ class Book(Document):
             ]}
         }
 
+        # Test embedded document in query
+        query = map_query({
+            'author': Author(name='JRR Tolkien', birthday=dt.datetime(1892, 1, 3))
+        }, book_fields)
+        assert isinstance(query['a'], dict)
+        assert query == {
+            'a': {'name': 'JRR Tolkien', 'b': dt.datetime(1892, 1, 3)}
+        }
+
+        # Test document in query
+        editor = Editor(name='Allen & Unwin')
+        editor.id = ObjectId()
+        query = map_query({'editor': editor}, book_fields)
+        assert isinstance(query['e'], ObjectId)
+        assert query['e'] == editor.id
+
     def test_mix(self):
 
         @self.instance.register

diff --git a/umongo/query_mapper.py b/umongo/query_mapper.py
@@ -1,4 +1,6 @@
 from umongo.fields import ListField, EmbeddedField
+from umongo.document import DocumentImplementation
+from umongo.embedded_document import EmbeddedDocumentImplementation
 
 
 def map_entry(entry, fields):
@@ -44,4 +46,9 @@ def map_query(query, fields):
         return mapped_query
     if isinstance(query, (list, tuple)):
         return [map_query(x, fields) for x in query]
+    # Passing a Document only makes sense in a Reference, let's query on ObjectId
+    if isinstance(query, DocumentImplementation):
+        return query.pk
+    if isinstance(query, EmbeddedDocumentImplementation):
+        return query.to_mongo()
     return query