From 47a228c50219004fd6a74e8f6dffbe062bf04086 Mon Sep 17 00:00:00 2001
From: Thomas Genin <6623268+tgenin@users.noreply.github.com>
Date: Mon, 31 Mar 2025 15:09:47 +0200
Subject: [PATCH] Allow ModelDependencies.get to retrieve missing models with
 None value

---
 modelkit/core/model.py |  6 ++++--
 tests/test_core.py     | 29 +++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/modelkit/core/model.py b/modelkit/core/model.py
index d72aafbd..2afb5d64 100644
--- a/modelkit/core/model.py
+++ b/modelkit/core/model.py
@@ -70,8 +70,10 @@ def __setitem__(self, key: str, value: ModelDependency) -> None:
 
     def get(
         self, key: str, model_type: Optional[Type[ModelDependency]] = None
-    ) -> ModelDependency:
-        m = self.models[key]
+    ) -> ModelDependency | None:
+        m = self.models.get(key)
+        if m is None:
+            return None
         if model_type and not isinstance(m, model_type):
             raise ValueError(f"Model `{m}` is not an instance of {model_type}")
         return cast(ModelDependency, m)
diff --git a/tests/test_core.py b/tests/test_core.py
index 422a856c..366a8308 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -665,6 +665,35 @@ def _predict(self, item):
     lib.get("some_model_dep")
 
 
+def test_model_dependencies_get_None():
+    class SomeModel(Model):
+        CONFIGURATIONS = {"some_model": {}}
+
+        def _predict(self, item):
+            return item
+
+    class SomeModelDep(Model):
+        CONFIGURATIONS = {
+            "some_model_dep_0": {"model_dependencies": {"some_model": "some_model"}},
+            "some_model_dep_1": {"model_dependencies": {}},
+        }
+
+        def _load(self):
+            self.some_model = self.model_dependencies.get("some_model")
+
+        def _predict(self, item):
+            return item
+
+    lib = ModelLibrary(
+        models=[SomeModel, SomeModelDep],
+        required_models=["some_model_dep_0", "some_model_dep_1"],
+    )
+    some_model_dep_0 = lib.get("some_model_dep_0")
+    some_model_dep_1 = lib.get("some_model_dep_1")
+    assert type(some_model_dep_0.some_model) == SomeModel
+    assert some_model_dep_1.some_model is None
+
+
 def test_asset_dependencies():
     """
     Test that MyAsset can have model dependencies of any kind