Fix mixin management in BaseInstance

Scille · Jul 20, 2020 · 5aed75f · 5aed75f
1 parent 36e77a9
commit 5aed75f
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/umongo/instance.py b/umongo/instance.py
@@ -1,6 +1,7 @@
 from .exceptions import (
     NotRegisteredDocumentError, AlreadyRegisteredDocumentError, NoDBDefinedError)
 from .document import DocumentTemplate
+from .embedded_document import EmbeddedDocumentTemplate
 from .template import get_template
 
 
@@ -34,6 +35,7 @@ def __init__(self, templates=()):
         self.builder = self.BUILDER_CLS(self)
         self._doc_lookup = {}
         self._embedded_lookup = {}
+        self._mixin_lookup = {}
         for template in templates:
             self.register(template)
 
@@ -99,8 +101,10 @@ class MyDoc(Document):
         template = get_template(template)
         if issubclass(template, DocumentTemplate):
             implementation = self._register_doc(template)
-        else:  # EmbeddedDocumentTemplate
+        elif issubclass(template, EmbeddedDocumentTemplate):
             implementation = self._register_embedded_doc(template)
+        else:  # MixinDocument
+            implementation = self._register_mixin_doc(template)
         if as_attribute:
             setattr(self, implementation.__name__, implementation)
         return implementation
@@ -126,7 +130,7 @@ def _register_mixin_doc(self, template):
         if hasattr(self, implementation.__name__):
             raise AlreadyRegisteredDocumentError(
                 'MixinDocument `%s` already registered' % implementation.__name__)
-        self._embedded_lookup[implementation.__name__] = implementation
+        self._mixin_lookup[implementation.__name__] = implementation
         return implementation