Fixing custom embedded documents regression

docs/source/_includes/substitutions.rst

@@ -133,5 +133,5 @@
 .. |AnnotationBackend| replace:: :class:`AnnotationBackend <fiftyone.utils.annotations.AnnotationBackend>`
 .. |AnnotationBackendConfig| replace:: :class:`AnnotationBackendConfig <fiftyone.utils.annotations.AnnotationBackendConfig>`
 
-.. |EmbeddedDocument| replace:: :class:`EmbeddedDocument <fiftyone.core.odm.document.EmbeddedDocument>`
-.. |DynamicEmbeddedDocument| replace:: :class:`DynamicEmbeddedDocument <fiftyone.core.odm.document.DynamicEmbeddedDocument>`
+.. |EmbeddedDocument| replace:: :class:`EmbeddedDocument <fiftyone.core.odm.embedded_document.EmbeddedDocument>`
+.. |DynamicEmbeddedDocument| replace:: :class:`DynamicEmbeddedDocument <fiftyone.core.odm.embedded_document.DynamicEmbeddedDocument>`

docs/source/release-notes.rst

@@ -556,9 +556,8 @@ Core
   cases)
 - Added support for indexing into datasets using boolean arrays or view
   expressions via new `dataset[bool_array]` and `dataset[bool_expr]` syntaxes
-- Added support for registering custom
-  :class:`EmbeddedDocument <fiftyone.core.odm.document.EmbeddedDocument>`
-  classes that can be used to populate fields and embedded fields of datasets
+- Added support for registering custom |EmbeddedDocument| classes that can be
+  used to populate fields and embedded fields of datasets
 - Added support for importing and exporting `confidence` in YOLO formats
 - Added support for directly passing a `filename -> filepath` mapping dict to
   the `data_path` parameter to

fiftyone/core/odm/utils.py

@@ -10,6 +10,7 @@
 import json
 import numbers
 import six
+import sys
 
 from bson import json_util
 from bson.binary import Binary
@@ -24,7 +25,8 @@
 import fiftyone.core.media as fom
 import fiftyone.core.utils as fou
 
-foed = fou.lazy_import("fiftyone.core.odm.embedded_document")
+food = fou.lazy_import("fiftyone.core.odm.document")
+fooe = fou.lazy_import("fiftyone.core.odm.embedded_document")
 
 
 def serialize_value(value, extended=False):
@@ -97,8 +99,8 @@ def deserialize_value(value):
     """
     if isinstance(value, dict):
         if "_cls" in value:
-            # Serialized embedded document
-            _cls = getattr(fo, value["_cls"])
+            # Serialized document
+            _cls = _document_registry[value["_cls"]]
             return _cls.from_dict(value)
 
         if "$binary" in value:
@@ -211,7 +213,7 @@ def get_implied_field_kwargs(value):
     Returns:
         a field specification dict
     """
-    if isinstance(value, foed.BaseEmbeddedDocument):
+    if isinstance(value, fooe.BaseEmbeddedDocument):
         return {
             "ftype": fof.EmbeddedDocumentField,
             "embedded_doc_type": type(value),
@@ -313,7 +315,7 @@ def _get_list_value_type(value):
     if isinstance(value, ObjectId):
         return fof.ObjectIdField
 
-    if isinstance(value, foed.BaseEmbeddedDocument):
+    if isinstance(value, fooe.BaseEmbeddedDocument):
         return fof.EmbeddedDocumentField
 
     if isinstance(value, datetime):
@@ -410,3 +412,66 @@ def _merge_field_kwargs(fields_list):
                 kwargs["fields"].append(v)
 
     return kwargs
+
+
+class DocumentRegistry(object):
+    """A registry of
+    :class:`fiftyone.core.odm.document.MongoEngineBaseDocument` classes found
+    when importing data from the database.
+    """
+
+    def __init__(self):
+        self._cache = {}
+
+    def __repr__(self):
+        return repr(self._cache)
+
+    def __getitem__(self, name):
+        # Check cache first
+        cls = self._cache.get(name, None)
+        if cls is not None:
+            return cls
+
+        # Then fiftyone namespace
+        try:
+            cls = self._get_cls(fo, name)
+            self._cache[name] = cls
+            return cls
+        except AttributeError:
+            pass
+
+        # Then full module list
+        for module in sys.modules.values():
+            try:
+                cls = self._get_cls(module, name)
+                self._cache[name] = cls
+                return cls
+            except AttributeError:
+                pass
+
+        raise DocumentRegistryError(
+            "Could not locate document class '%s'.\n\nIf you are working with "
+            "a dataset that uses custom embedded documents, you must add them "
+            "to FiftyOne's module path. See "
+            "https://voxel51.com/docs/fiftyone/user_guide/using_datasets.html#custom-embedded-documents "
+            "for more information" % name
+        )
+
+    def _get_cls(self, module, name):
+        cls = getattr(module, name)
+
+        try:
+            assert issubclass(cls, food.MongoEngineBaseDocument)
+        except:
+            raise AttributeError
+
+        return cls
+
+
+class DocumentRegistryError(Exception):
+    """Error raised when an unknown document class is encountered."""
+
+    pass
+
+
+_document_registry = DocumentRegistry()

tests/unittests/dataset_tests.py

@@ -18,6 +18,7 @@
 
 import fiftyone as fo
 from fiftyone import ViewField as F
+import fiftyone.core.fields as fof
 import fiftyone.core.odm as foo
 
 from decorators import drop_datasets, skip_windows
@@ -2784,6 +2785,50 @@ def test_delete_video_detections_labels(self):
         self.assertEqual(num_labels_after, num_labels - num_selected)
 
 
+class CustomEmbeddedDocumentTests(unittest.TestCase):
+    @drop_datasets
+    def test_custom_embedded_documents(self):
+        sample = fo.Sample(
+            filepath="/path/to/image.png",
+            camera_info=_CameraInfo(
+                camera_id="123456789",
+                quality=99.0,
+            ),
+            weather=fo.Classification(
+                label="sunny",
+                confidence=0.95,
+                metadata=_LabelMetadata(
+                    model_name="resnet50",
+                    description="A dynamic field",
+                ),
+            ),
+        )
+
+        dataset = fo.Dataset()
+        dataset.add_sample(sample)
+
+        self.assertIsInstance(sample.camera_info, _CameraInfo)
+        self.assertIsInstance(sample.weather.metadata, _LabelMetadata)
+
+        view = dataset.limit(1)
+        sample_view = view.first()
+
+        self.assertIsInstance(sample_view.camera_info, _CameraInfo)
+        self.assertIsInstance(sample_view.weather.metadata, _LabelMetadata)
+
+
+class _CameraInfo(foo.EmbeddedDocument):
+    camera_id = fof.StringField(required=True)
+    quality = fof.FloatField()
+    description = fof.StringField()
+
+
+class _LabelMetadata(foo.DynamicEmbeddedDocument):
+    created_at = fof.DateTimeField(default=datetime.utcnow)
+
+    model_name = fof.StringField()
+
+
 if __name__ == "__main__":
     fo.config.show_progress_bars = False
     unittest.main(verbosity=2)