Resolve model engine names through the shared registry (item 5, step 2)

src/Auto3D/model_factory.py

@@ -22,6 +22,7 @@
     ANI2xtAdapter,
     CustomModelAdapter,
 )
+from Auto3D.registry import Registry
 
 if TYPE_CHECKING:
     # Annotation-only. Every signature here promises the CONTRACT
@@ -57,13 +58,27 @@ class ModelFactory:
         >>> ModelFactory.clear_cache()
     """
 
-    # Built-in (non-aimnet) engines kept for back-compat; keys are exactly the
-    # members of BUILTIN_ANI_MODELS.
-    _adapters: dict[str, type[BaseModelAdapter]] = {
-        MODEL_ANI2XT.upper(): ANI2xtAdapter,
-        MODEL_ANI2X.upper(): ANI2xAdapter,
-    }
-    assert set(_adapters) == set(BUILTIN_ANI_MODELS)
+    #: Every engine name a user may pass, in the order `auto3d models list`
+    #: shows them. The value is the adapter class for engines built from one
+    #: locally, and ``None`` for names resolved through the aimnet registry --
+    #: which is the distinction that made this three separate lists before.
+    #:
+    #: `_adapters` held only the two ANI engines, `available_models()` restated
+    #: all six as a hand-written literal derived from nothing, and
+    #: `cli/commands/models.py`'s `ENGINE_INFO` held six display entries keyed by
+    #: the same names. Nothing connected them; they agreed by hand. Two of the
+    #: three are collapsed here. `ENGINE_INFO` is the remaining one, and moving
+    #: it into `info=` is the follow-up -- it belongs with the CLI changes rather
+    #: than with this factory's.
+    _engines: Registry[type[BaseModelAdapter] | None] = Registry(
+        "optimizing engine", case_insensitive=True
+    )
+    _engines.register(MODEL_AIMNET, None)
+    _engines.register("aimnet2-2025", None)
+    _engines.register("aimnet2-nse", None)
+    _engines.register("aimnet2-pd", None)
+    _engines.register(MODEL_ANI2X, ANI2xAdapter)
+    _engines.register(MODEL_ANI2XT, ANI2xtAdapter)
 
     # Model instance cache: key = (name, device_str, compile_model)
     _cache: dict[tuple[str, str, bool], ModelAdapter] = {}
@@ -141,12 +156,14 @@ def create(
         #    name first; that helper is gone -- the Hessian path takes a
         #    ModelAdapter now -- so this factory is the only name resolver
         #    left, and the reason above stands on its own.)
-        if name_upper in cls._adapters:
+        if name_upper in cls._engines and cls._engines.resolve(name_upper) is not None:
             cache_key = (name_upper, str(device), compile_model)
             if use_cache and cache_key in cls._cache:
                 return cls._cache[cache_key]
             try:
-                adapter = cls._adapters[name_upper](device, compile_model=compile_model)
+                adapter_cls = cls._engines.resolve(name_upper)
+                assert adapter_cls is not None  # guarded by the branch condition
+                adapter = adapter_cls(device, compile_model=compile_model)
             except ImportError as exc:
                 # Only translate the *absence of torchani itself*. A broken
                 # torchani whose own transitive import fails names a different
@@ -185,15 +202,25 @@ def create(
 
     @classmethod
     def available_models(cls) -> list[str]:
-        """Return list of registered model names."""
-        return [
-            MODEL_AIMNET,
-            "aimnet2-2025",
-            "aimnet2-nse",
-            "aimnet2-pd",
-            MODEL_ANI2X,
-            MODEL_ANI2XT,
-        ]
+        """Registered engine names, in declaration order.
+
+        Was a hand-written literal that restated the registry's contents and was
+        derived from nothing -- so adding an engine and forgetting this list left
+        it working but undiscoverable.
+        """
+        return cls._engines.available()
+
+
+# The engines built from a local adapter class must be exactly
+# BUILTIN_ANI_MODELS. This was an assert over a hand-maintained dict; it is now
+# derived from the registry, so the two cannot drift apart. At module scope
+# because a comprehension inside a class body cannot see that class's own
+# attributes.
+assert {
+    name.upper()
+    for name in ModelFactory._engines.available()
+    if ModelFactory._engines.resolve(name) is not None
+} == set(BUILTIN_ANI_MODELS)
 
 
 def create_model(

tests/test_cli_exit_codes.py

@@ -345,7 +345,11 @@ def _broken_adapter(device, compile_model=False):
             "No module named 'some_transitive_dep'", name="some_transitive_dep"
         )
 
-    monkeypatch.setitem(mf.ModelFactory._adapters, "ANI2X", _broken_adapter)
+    monkeypatch.setitem(
+        mf.ModelFactory._engines._entries,
+        "ANI2x",
+        mf.ModelFactory._engines.entry("ANI2x").__class__(name="ANI2x", value=_broken_adapter),
+    )
     monkeypatch.setattr(mf.ModelFactory, "_cache", {})
 
     result = runner.invoke(app, ["models", "test", "ANI2x", "--no-gpu"])

tests/test_model_factory.py

@@ -35,7 +35,10 @@ def test_the_factory_promises_the_contract_not_the_base_class():
     assert inspect.get_annotations(ModelFactory.create.__func__)["return"] == "ModelAdapter"
     # BaseModelAdapter survives in exactly one type position: a registry of
     # Auto3D's OWN adapter classes, which really is the concrete base.
-    assert inspect.get_annotations(ModelFactory)["_adapters"] == "dict[str, type[BaseModelAdapter]]"
+    assert (
+        inspect.get_annotations(ModelFactory)["_engines"]
+        == "Registry[type[BaseModelAdapter] | None]"
+    )
     # ...and it is no longer an incidental runtime re-export of this module.
     assert not hasattr(model_factory, "BaseModelAdapter")
 
@@ -51,8 +54,17 @@ def test_registry_is_populated(self):
         assert "ANI2x" in models
         assert "ANI2xt" in models
         # AIMNET is NOT a hard-coded adapter key anymore: only ANI engines are.
-        assert set(ModelFactory._adapters) == {"ANI2X", "ANI2XT"}
-        assert "AIMNET" not in ModelFactory._adapters
+        # The registry now holds every user-facing engine name; the ones built
+        # from a local adapter class are those whose entry carries one. AIMNET
+        # is registered (it is a name a user may pass) but resolves to None,
+        # because it is looked up through the aimnet registry instead.
+        built_locally = {
+            n.upper()
+            for n in ModelFactory._engines.available()
+            if ModelFactory._engines.resolve(n) is not None
+        }
+        assert built_locally == {"ANI2X", "ANI2XT"}
+        assert ModelFactory._engines.resolve("AIMNET") is None
 
     def test_create_unknown_model_raises_error(self, monkeypatch):
         """Unknown non-path names no longer raise a ValueError up front; they
@@ -93,11 +105,10 @@ def __init__(self, model_name, device, **kw):
             assert captured["aimnet"] == "aimnet2"
 
         # ANI engines resolve case-insensitively to their adapter class.
-        assert (
-            model_factory.ModelFactory._adapters["ani2x".upper()]
-            is model_factory.ModelFactory._adapters["ANI2X"]
-        )
-        assert "ANI2XT" in model_factory.ModelFactory._adapters
+        assert model_factory.ModelFactory._engines.resolve(
+            "ani2x"
+        ) is model_factory.ModelFactory._engines.resolve("ANI2X")
+        assert "ANI2XT" in model_factory.ModelFactory._engines
 
     @pytest.mark.slow
     def test_create_aimnet_returns_aimnet2_adapter(self, aimnet_model):
@@ -364,7 +375,13 @@ class _FakeANI2xtAdapter:
         def __init__(self, device, **kw):
             captured["built_in"] = True
 
-    monkeypatch.setitem(model_factory.ModelFactory._adapters, "ANI2XT", _FakeANI2xtAdapter)
+    monkeypatch.setitem(
+        model_factory.ModelFactory._engines._entries,
+        "ANI2xt",
+        model_factory.ModelFactory._engines.entry("ANI2xt").__class__(
+            name="ANI2xt", value=_FakeANI2xtAdapter
+        ),
+    )
     model_factory.ModelFactory.clear_cache()
 
     result = model_factory.create_model("ANI2xt", torch.device("cpu"), use_cache=False)