Merge c9a8664 into 2f3322c

tests/test_document.py

@@ -351,7 +351,7 @@ class Cousin(Uncle):
         assert Child.opts.offspring == set()
         assert Cousin.opts.offspring == set()
 
-    def test_instanciate_template(self):
+    def test_instantiate_template(self):
 
         class Doc(Document):
             pass
@@ -441,6 +441,68 @@ class User(Document):
         with pytest.raises(exceptions.AlreadyCreatedError):
             john.update({'primary_key': ObjectId()})
 
+    def test_mixin(self):
+
+        class PMixin:
+            pm = fields.IntField()
+
+        class CMixin:
+            cm = fields.IntField()
+
+        @self.instance.register
+        class Parent(PMixin, Document):
+            p = fields.StrField()
+
+            class Meta:
+                allow_inheritance = True
+
+        @self.instance.register
+        class Child(CMixin, Parent):
+            c = fields.StrField()
+
+        assert set(Parent.schema.fields.keys()) == {'id', 'p', 'pm'}
+        assert set(Child.schema.fields.keys()) == {'id', 'cls', 'p', 'pm', 'c', 'cm'}
+
+        parent_data = {'p': 'parent', 'pm': 42}
+        child_data = {'c': 'Child', 'cm': 12, **parent_data}
+
+        assert Parent(**parent_data).dump() == parent_data
+        assert Child(**child_data).dump() == {**child_data, 'cls': 'Child'}
+
+    def test_mixin_override(self):
+
+        class PMixin:
+            pm = fields.IntField()
+
+        class CMixin:
+            cm = fields.IntField()
+
+        @self.instance.register
+        class Parent(PMixin, Document):
+            p = fields.StrField()
+            pm = fields.IntField(validate=ma.validate.Range(0, 5))
+
+            class Meta:
+                allow_inheritance = True
+
+        @self.instance.register
+        class Child(CMixin, Parent):
+            c = fields.StrField()
+            cm = fields.IntField(validate=ma.validate.Range(0, 5))
+
+        assert set(Parent.schema.fields.keys()) == {'id', 'p', 'pm'}
+        assert set(Child.schema.fields.keys()) == {'id', 'cls', 'p', 'pm', 'c', 'cm'}
+
+        parent_data = {'p': 'parent', 'pm': 42}
+        child_data = {'c': 'Child', 'cm': 12, **parent_data}
+
+        with pytest.raises(ma.ValidationError) as exc:
+            Parent(**parent_data)
+        assert set(exc.value.messages.keys()) == {'pm'}
+        with pytest.raises(ma.ValidationError) as exc:
+            Child(**child_data)
+        assert set(exc.value.messages.keys()) == {'pm', 'cm'}
+
 
 class TestConfig(BaseTest):
 

tests/test_embedded_document.py

@@ -441,3 +441,65 @@ class Parent(EmbeddedDocument):
         assert jane.name == john.name
         assert jane.child == john.child
         assert jane.child is not john.child
+
+    def test_mixin(self):
+
+        class PMixin:
+            pm = fields.IntField()
+
+        class CMixin:
+            cm = fields.IntField()
+
+        @self.instance.register
+        class Parent(PMixin, EmbeddedDocument):
+            p = fields.StrField()
+
+            class Meta:
+                allow_inheritance = True
+
+        @self.instance.register
+        class Child(CMixin, Parent):
+            c = fields.StrField()
+
+        assert set(Parent.schema.fields.keys()) == {'p', 'pm'}
+        assert set(Child.schema.fields.keys()) == {'cls', 'p', 'pm', 'c', 'cm'}
+
+        parent_data = {'p': 'parent', 'pm': 42}
+        child_data = {'c': 'Child', 'cm': 12, **parent_data}
+
+        assert Parent(**parent_data).dump() == parent_data
+        assert Child(**child_data).dump() == {**child_data, 'cls': 'Child'}
+
+    def test_mixin_override(self):
+
+        class PMixin:
+            pm = fields.IntField()
+
+        class CMixin:
+            cm = fields.IntField()
+
+        @self.instance.register
+        class Parent(PMixin, EmbeddedDocument):
+            p = fields.StrField()
+            pm = fields.IntField(validate=ma.validate.Range(0, 5))
+
+            class Meta:
+                allow_inheritance = True
+
+        @self.instance.register
+        class Child(CMixin, Parent):
+            c = fields.StrField()
+            cm = fields.IntField(validate=ma.validate.Range(0, 5))
+
+        assert set(Parent.schema.fields.keys()) == {'p', 'pm'}
+        assert set(Child.schema.fields.keys()) == {'cls', 'p', 'pm', 'c', 'cm'}
+
+        parent_data = {'p': 'parent', 'pm': 42}
+        child_data = {'c': 'Child', 'cm': 12, **parent_data}
+
+        with pytest.raises(ma.ValidationError) as exc:
+            Parent(**parent_data)
+        assert set(exc.value.messages.keys()) == {'pm'}
+        with pytest.raises(ma.ValidationError) as exc:
+            Child(**child_data)
+        assert set(exc.value.messages.keys()) == {'pm', 'cm'}

umongo/builder.py

@@ -84,17 +84,23 @@ def _collect_schema_attrs(template):
     schema_fields = {}
     schema_non_fields = {}
     nmspc = {}
-    for key, item in template.__dict__.items():
-        if hasattr(item, '__marshmallow_hook__'):
-            # Decorated special functions (e.g. `post_load`)
-            schema_non_fields[key] = item
-        elif isinstance(item, ma.fields.Field):
-            # Given the fields provided by the template are going to be
-            # customized in the implementation, we copy them to avoid
-            # overwriting if two implementations are created
-            schema_fields[key] = copy(item)
-        else:
-            nmspc[key] = item
+    # Collect schema attributes up the MRO to get attributes from Mixins
+    for tmpl in inspect.getmro(template):
+        # Stop when reaching Doc or EmbeddedDoc template as this is already
+        # covered by document inheritance
+        if tmpl != template and issubclass(tmpl, Template):
+            break
+        for key, item in tmpl.__dict__.items():
+            if hasattr(item, '__marshmallow_hook__'):
+                # Decorated special functions (e.g. `post_load`)
+                schema_non_fields.setdefault(key, item)
+            elif isinstance(item, ma.fields.Field):
+                # Given the fields provided by the template are going to be
+                # customized in the implementation, we copy them to avoid
+                # overwriting if two implementations are created
+                schema_fields.setdefault(key, copy(item))
+            else:
+                nmspc.setdefault(key, item)
     return nmspc, schema_fields, schema_non_fields