Repository: Scille/umongo
Pull request #261 (2020-04-28 22:55:29)

Add ExposeMissing context manager
- Use Python 3.7's context variable feature to create a context manager in which a document returns missing for missing data rather than None
- Use this context manager in schema_from_umongo_get_attribute

This makes the Document responsible for outputing missing rather than the schema having to guess. Also the schema works with any type of data, not only a Document instance (with other data, the context manger is no-op).

I hoped to remove the whole `mongo_world` management thanks to this, but there are other issues with data from Mongo vs. Document.

The benefit is not as great as I expected, but I still like this implementation.

coveralls: 
[![Coverage Status](https://coveralls.io/builds/30416147/badge)](https://coveralls.io/builds/30416147)

Coverage decreased (-20.9%) to 74.852% when pulling **021a3b3e3735a96607cc4dce9c563f1766eb8309 on expose_missing** into **8a0aaf9390ed57b9bcdc8057af241be58110f5de on master**.

coveralls: 
[![Coverage Status](https://coveralls.io/builds/30416147/badge)](https://coveralls.io/builds/30416147)

Coverage decreased (-20.9%) to 74.852% when pulling **021a3b3e3735a96607cc4dce9c563f1766eb8309 on expose_missing** into **8a0aaf9390ed57b9bcdc8057af241be58110f5de on master**.

lafrech: This implementation is totally sub-optimal. It misses a nice point of the context variable, which is that we could put the context manager at the dump method level rather than the get_attribute level. This is another reason to like the context_var approach better than the current one.
lafrech: Rebased on master.

- Now uses the context manager in schema dump method.
- Removes now useless `schema_from_umongo_get_attribute`.
- Exposes (`SchemaFomUmongo` renamed as) `RemoveMissingSchema`.
- Makes `RemoveMissingSchema` an opt-in. Easy to achieve by passing it to custom `BaseDocument` and `BaseEmbeddedDocument`.

Patch:
diff --git a/docs/userguide.rst b/docs/userguide.rst
index eb06332..93a303d 100644
--- a/docs/userguide.rst
+++ b/docs/userguide.rst
@@ -569,7 +569,7 @@ from scratch... almost:
     >>> MassiveDuckSchema(marshmallow.Schema):
     ...     breeds = marshmallow.fields.List(marshmallow.fields.Nested(MassiveBreedSchema))
 
-.. note:: A custom marshmallow schema :class:`umongo.marshmallow_bonus.SchemaFromUmongo`
+.. note:: A custom marshmallow schema :class:`umongo.schema.RemoveMissingSchema`
     can be used instead of regular :class:`marshmallow.Schema` to skip missing fields
     when dumping a :class:`umongo.Document` object.
 
diff --git a/umongo/__init__.py b/umongo/__init__.py
index 2557191..caa60a6 100644
--- a/umongo/__init__.py
+++ b/umongo/__init__.py
@@ -29,6 +29,7 @@ from .schema import Schema
 from .data_objects import Reference
 from .embedded_document import EmbeddedDocument
 from .mixin import MixinDocument
+from .expose_missing import ExposeMissing
 from .i18n import set_gettext
 
 
@@ -52,6 +53,7 @@ __all__ = (
     'validates_schema',
     'EmbeddedDocument',
     'MixinDocument',
+    'ExposeMissing',
 
     'UMongoError',
     'ValidationError',
diff --git a/umongo/embedded_document.py b/umongo/embedded_document.py
index f052eb9..529bb37 100644
--- a/umongo/embedded_document.py
+++ b/umongo/embedded_document.py
@@ -3,6 +3,7 @@ import marshmallow as ma
 
 from .template import Implementation, Template
 from .data_objects import BaseDataObject
+from .expose_missing import EXPOSE_MISSING
 from .exceptions import AbstractDocumentError
 
 
@@ -159,7 +160,7 @@ class EmbeddedDocumentImplementation(Implementation, BaseDataObject):
 
     def __getitem__(self, name):
         value = self._data.get(name)
-        return value if value is not ma.missing else None
+        return None if value is ma.missing and not EXPOSE_MISSING.get() else value
 
     def __delitem__(self, name):
         self._data.delete(name)
@@ -182,7 +183,7 @@ class EmbeddedDocumentImplementation(Implementation, BaseDataObject):
         if name[:2] == name[-2:] == '__':
             raise AttributeError(name)
         value = self._data.get(name, to_raise=AttributeError)
-        return value if value is not ma.missing else None
+        return None if value is ma.missing and not EXPOSE_MISSING.get() else value
 
     def __delattr__(self, name):
         if not self.__real_attributes:
diff --git a/umongo/expose_missing.py b/umongo/expose_missing.py
new file mode 100644
index 0000000..4454e35
--- /dev/null
+++ b/umongo/expose_missing.py
@@ -0,0 +1,26 @@
+"""Expose missing context variable
+
+Allows the user to let umongo document return missing rather than None for
+empty fields.
+"""
+
+
+from contextvars import ContextVar
+from contextlib import AbstractContextManager
+
+EXPOSE_MISSING = ContextVar("expose_missing", default=False)
+
+
+class ExposeMissing(AbstractContextManager):
+    """Let Document expose missing values rather than returning None
+
+    By default, getting a document item returns None if the value is missing.
+    Inside this context manager, the missing singleton is returned. This can
+    be useful is cases where the user want to distinguish between None and
+    missing value.
+    """
+    def __enter__(self):
+        self.token = EXPOSE_MISSING.set(True)
+
+    def __exit__(self, *args, **kwargs):
+        EXPOSE_MISSING.reset(self.token)
diff --git a/umongo/schema.py b/umongo/schema.py
index 3a4ad6f..f53a4ea 100644
--- a/umongo/schema.py
+++ b/umongo/schema.py
@@ -3,45 +3,22 @@ import marshmallow as ma
 
 from .abstract import BaseSchema
 from .i18n import gettext as _
-
+from .expose_missing import ExposeMissing
 
 __all__ = (
     'Schema',
-    'schema_from_umongo_get_attribute',
-    'SchemaFromUmongo',
+    'RemoveMissingSchema',
 )
 
 
-def schema_from_umongo_get_attribute(self, obj, attr, default):
-    """
-    Overwrite default `Schema.get_attribute` method by this one to access
-        umongo missing fields instead of returning `None`.
-
-    example::
-
-        class MySchema(marshsmallow.Schema):
-            get_attribute = schema_from_umongo_get_attribute
-
-            # Define the rest of your schema
-            ...
-
+class RemoveMissingSchema(ma.Schema):
     """
-    ret = ma.Schema.get_attribute(self, obj, attr, default)
-    if ret is None and ret is not default and attr in obj.schema.fields:
-        raw_ret = obj._data.get(attr)
-        return default if raw_ret is ma.missing else raw_ret
-    return ret
-
-
-class SchemaFromUmongo(ma.Schema):
-    """
-    Custom :class:`marshmallow.Schema` subclass providing unknown fields
-    checking and custom get_attribute for umongo documents.
-
-    .. note: It is not mandatory to use this schema with umongo document.
-        This is just a helper providing usefull behaviors.
+    Custom :class:`marshmallow.Schema` subclass returning missing rather than
+    None for missing fields in umongo :class:`umongo.Document`s.
     """
-    get_attribute = schema_from_umongo_get_attribute
+    def dump(self, *args, **kwargs):
+        with ExposeMissing():
+            return super().dump(*args, **kwargs)
 
 
 class Schema(BaseSchema):
@@ -67,10 +44,6 @@ class Schema(BaseSchema):
             for name, field in self.fields.items()
         }
         name = 'Marshmallow%s' % type(self).__name__
-        # By default OO world returns `missing` fields as `None`,
-        # disable this behavior here to let marshmallow deal with it
-        if not mongo_world:
-            nmspc['get_attribute'] = schema_from_umongo_get_attribute
         m_schema = type(name, (self.MA_BASE_SCHEMA_CLS, ), nmspc)
         # Add i18n support to the schema
         # We can't use I18nErrorDict here because __getitem__ is not called


Test patch:
diff --git a/tests/test_document.py b/tests/test_document.py
index 1844b2d..0fac665 100644
--- a/tests/test_document.py
+++ b/tests/test_document.py
@@ -8,9 +8,8 @@ import marshmallow as ma
 
 from umongo import (
     Document, EmbeddedDocument, MixinDocument,
-    Schema, fields, exceptions, post_dump, pre_load, validates_schema
+    Schema, fields, exceptions, post_dump, pre_load, validates_schema, ExposeMissing
 )
-
 from .common import BaseTest
 
 
@@ -434,6 +433,14 @@ class TestDocument(BaseTest):
         with pytest.raises(exceptions.AlreadyCreatedError):
             john.update({'primary_key': ObjectId()})
 
+    def test_expose_missing(self):
+        john = self.Student(name='John Doe')
+        assert john.name == 'John Doe'
+        assert john.birthday is None
+        with ExposeMissing():
+            assert john.name == 'John Doe'
+            assert john.birthday is ma.missing
+
     def test_mixin(self):
 
         @self.instance.register
diff --git a/tests/test_embedded_document.py b/tests/test_embedded_document.py
index 9937c49..52c36f9 100644
--- a/tests/test_embedded_document.py
+++ b/tests/test_embedded_document.py
@@ -5,7 +5,7 @@ import pytest
 import marshmallow as ma
 
 from umongo.data_proxy import data_proxy_factory
-from umongo import Document, EmbeddedDocument, MixinDocument, fields, exceptions
+from umongo import Document, EmbeddedDocument, MixinDocument, fields, exceptions, ExposeMissing
 
 from .common import BaseTest
 
@@ -414,6 +414,26 @@ class TestEmbeddedDocument(BaseTest):
         assert jane.child == john.child
         assert jane.child is not john.child
 
+    def test_expose_missing(self):
+        @self.instance.register
+        class Child(EmbeddedDocument):
+            name = fields.StrField()
+            age = fields.IntField()
+
+        @self.instance.register
+        class Parent(EmbeddedDocument):
+            name = fields.StrField()
+            child = fields.EmbeddedField(Child)
+
+        parent = Parent(**{'child': {'age': 42}})
+        assert parent.name is None
+        assert parent.child.name is None
+        assert parent.child.age == 42
+        with ExposeMissing():
+            assert parent.name is ma.missing
+            assert parent.child.name is ma.missing
+            assert parent.child.age == 42
+
     def test_mixin(self):
 
         @self.instance.register
diff --git a/tests/test_marshmallow.py b/tests/test_marshmallow.py
index f03efed..26df886 100644
--- a/tests/test_marshmallow.py
+++ b/tests/test_marshmallow.py
@@ -9,7 +9,7 @@ import marshmallow as ma
 from umongo import Document, EmbeddedDocument, fields, set_gettext, validate, missing
 from umongo import marshmallow_bonus as ma_bonus_fields, Schema
 from umongo.abstract import BaseField
-from umongo.schema import schema_from_umongo_get_attribute, SchemaFromUmongo
+from umongo.schema import RemoveMissingSchema
 
 from .common import BaseTest
 
@@ -52,10 +52,16 @@ class TestMarshmallow(BaseTest):
         class MyDocument(Document):
             MA_BASE_SCHEMA_CLS = ExcludeBaseSchema
 
+            class Meta:
+                abstract = True
+
         @self.instance.register
         class MyEmbeddedDocument(EmbeddedDocument):
             MA_BASE_SCHEMA_CLS = ExcludeBaseSchema
 
+            class Meta:
+                abstract = True
+
         # Now, all our objects will generate "exclude" marshmallow schemas
         @self.instance.register
         class Accessory(MyEmbeddedDocument):
@@ -377,7 +383,7 @@ class TestMarshmallow(BaseTest):
                 schema.load({"id": invalid_id})
             assert exc.value.messages == {"id": ["Invalid ObjectId."]}
 
-    def test_marshmallow_schema_helpers(self):
+    def test_marshmallow_remove_missing_schema(self):
 
         @self.instance.register
         class Doc(Document):
@@ -387,32 +393,80 @@ class TestMarshmallow(BaseTest):
             def prop(self):
                 return "I'm a property !"
 
-        class VanillaSchema(ma.Schema):
+        class VanillaDocSchema(ma.Schema):
             a = ma.fields.Int()
 
-        class CustomGetAttributeSchema(VanillaSchema):
-            get_attribute = schema_from_umongo_get_attribute
+        class RemoveMissingDocSchema(RemoveMissingSchema):
+            a = ma.fields.Int()
 
-        data = VanillaSchema().dump(Doc())
+        data = VanillaDocSchema().dump(Doc())
         assert data == {'a': None}
 
-        data = CustomGetAttributeSchema().dump(Doc())
+        data = RemoveMissingDocSchema().dump(Doc())
         assert data == {}
 
-        data = CustomGetAttributeSchema().dump(Doc(a=1))
+        data = RemoveMissingDocSchema().dump(Doc(a=1))
         assert data == {'a': 1}
 
-        class MySchemaFromUmongo(SchemaFromUmongo):
-            a = ma.fields.Int()
-            prop = ma.fields.String(dump_only=True)
+    @pytest.mark.parametrize("base_schema", (ma.Schema, RemoveMissingSchema))
+    def test_marshmallow_remove_missing_schema_as_base_schema(self, base_schema):
+        """Test RemoveMissingSchema used as base marshmallow Schema"""
 
-        with pytest.raises(ma.ValidationError) as excinfo:
-            MySchemaFromUmongo().load({'a': 1, 'dummy': 2})
-        assert excinfo.value.messages == {'dummy': ['Unknown field.']}
+        # Typically, we'll use it in all our schemas, so let's define base
+        # Document and EmbeddedDocument classes using this base schema class
+        @self.instance.register
+        class MyDocument(Document):
+            MA_BASE_SCHEMA_CLS = base_schema
 
-        with pytest.raises(ma.ValidationError) as excinfo:
-            MySchemaFromUmongo().load({'a': 1, 'prop': '2'})
-        assert excinfo.value.messages == {'prop': ['Unknown field.']}
+            class Meta:
+                abstract = True
+
+        @self.instance.register
+        class MyEmbeddedDocument(EmbeddedDocument):
+            MA_BASE_SCHEMA_CLS = base_schema
+
+            class Meta:
+                abstract = True
+
+        @self.instance.register
+        class Accessory(MyEmbeddedDocument):
+            brief = fields.StrField()
+            value = fields.IntField()
+
+        @self.instance.register
+        class Bag(MyDocument):
+            item = fields.EmbeddedField(Accessory)
+            content = fields.ListField(fields.EmbeddedField(Accessory))
+
+        data = {
+            'item': {'brief': 'sportbag'},
+            'content': [
+                {'brief': 'cellphone'},
+                {'brief': 'lighter'}]
+        }
+        dump = {
+            'id': None,
+            'content': [
+                {'brief': 'cellphone', 'value': None},
+                {'brief': 'lighter', 'value': None}
+            ],
+            'item': {'brief': 'sportbag', 'value': None}
+        }
+        remove_missing_dump = {
+            'item': {'brief': 'sportbag'},
+            'content': [
+                {'brief': 'cellphone'},
+                {'brief': 'lighter'}
+            ]
+        }
+        expected_dump = {
+            ma.Schema: dump,
+            RemoveMissingSchema: remove_missing_dump,
+        }[base_schema]
+
+        bag = Bag(**data)
+        ma_schema = Bag.schema.as_marshmallow_schema()
+        assert ma_schema().dump(bag) == expected_dump
 
     def test_marshmallow_access_custom_attributes(self):