refactor: implement registries as mutable mappings

lightx2v/utils/registry_factory.py

@@ -1,10 +1,11 @@
+from collections.abc import MutableMapping
+
 from lightx2v_platform.registry_factory import PLATFORM_ATTN_WEIGHT_REGISTER, PLATFORM_LAYERNORM_WEIGHT_REGISTER, PLATFORM_MM_WEIGHT_REGISTER, PLATFORM_RMS_WEIGHT_REGISTER, PLATFORM_ROPE_REGISTER
 
 
-class Register(dict):
+class Register(MutableMapping):
     def __init__(self, *args, **kwargs):
-        super(Register, self).__init__(*args, **kwargs)
-        self._dict = {}
+        self._dict = dict(*args, **kwargs)
 
     def __call__(self, target_or_name):
         if callable(target_or_name):
@@ -31,6 +32,15 @@ def __setitem__(self, key, value):
     def __getitem__(self, key):
         return self._dict[key]
 
+    def __delitem__(self, key):
+        del self._dict[key]
+
+    def __iter__(self):
+        return iter(self._dict)
+
+    def __len__(self):
+        return len(self._dict)
+
     def __contains__(self, key):
         return key in self._dict
 

lightx2v_platform/base/intel_xpu.py

@@ -60,4 +60,4 @@ def get_device() -> str:
 
 
 # Register alias "xpu" for backward compatibility
-PLATFORM_DEVICE_REGISTER._dict["xpu"] = IntelXpuDevice
+PLATFORM_DEVICE_REGISTER["xpu"] = IntelXpuDevice

lightx2v_platform/registry_factory.py

@@ -1,7 +1,9 @@
-class Register(dict):
+from collections.abc import MutableMapping
+
+
+class Register(MutableMapping):
     def __init__(self, *args, **kwargs):
-        super(Register, self).__init__(*args, **kwargs)
-        self._dict = {}
+        self._dict = dict(*args, **kwargs)
 
     def __call__(self, target_or_name):
         if callable(target_or_name):
@@ -28,6 +30,15 @@ def __setitem__(self, key, value):
     def __getitem__(self, key):
         return self._dict[key]
 
+    def __delitem__(self, key):
+        del self._dict[key]
+
+    def __iter__(self):
+        return iter(self._dict)
+
+    def __len__(self):
+        return len(self._dict)
+
     def __contains__(self, key):
         return key in self._dict
 

lightx2v_train/lightx2v_train/utils/registry.py

@@ -1,7 +1,9 @@
-class Register(dict):
+from collections.abc import MutableMapping
+
+
+class Register(MutableMapping):
     def __init__(self, *args, **kwargs):
-        super(Register, self).__init__(*args, **kwargs)
-        self._dict = {}
+        self._dict = dict(*args, **kwargs)
 
     def __call__(self, target_or_name):
         if callable(target_or_name):
@@ -28,6 +30,15 @@ def __setitem__(self, key, value):
     def __getitem__(self, key):
         return self._dict[key]
 
+    def __delitem__(self, key):
+        del self._dict[key]
+
+    def __iter__(self):
+        return iter(self._dict)
+
+    def __len__(self):
+        return len(self._dict)
+
     def __contains__(self, key):
         return key in self._dict
 

test_cases/test_registry_factory.py

@@ -0,0 +1,147 @@
+import importlib.util
+import sys
+import types
+import unittest
+from collections.abc import MutableMapping
+from pathlib import Path
+from unittest.mock import patch
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_module(module_name, path):
+    spec = importlib.util.spec_from_file_location(module_name, path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Cannot load module {module_name} from {path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def first_target():
+    return "first"
+
+
+def second_target():
+    return "second"
+
+
+platform_registry = load_module("_platform_registry_for_test", REPO_ROOT / "lightx2v_platform" / "registry_factory.py")
+
+platform_package = types.ModuleType("lightx2v_platform")
+platform_package.registry_factory = platform_registry
+with patch.dict(
+    sys.modules,
+    {
+        "lightx2v_platform": platform_package,
+        "lightx2v_platform.registry_factory": platform_registry,
+    },
+):
+    runtime_registry = load_module("_runtime_registry_for_test", REPO_ROOT / "lightx2v" / "utils" / "registry_factory.py")
+
+training_registry = load_module(
+    "_training_registry_for_test",
+    REPO_ROOT / "lightx2v_train" / "lightx2v_train" / "utils" / "registry.py",
+)
+
+
+REGISTER_CLASSES = {
+    "runtime": runtime_registry.Register,
+    "platform": platform_registry.Register,
+    "training": training_registry.Register,
+}
+
+
+class RegisterTest(unittest.TestCase):
+    def test_implements_mutable_mapping(self):
+        for name, register_cls in REGISTER_CLASSES.items():
+            with self.subTest(register=name):
+                self.assertTrue(issubclass(register_cls, MutableMapping))
+                self.assertFalse(issubclass(register_cls, dict))
+
+    def test_initial_mapping_uses_registry_storage(self):
+        def initial_target():
+            return "initial"
+
+        def keyword_target():
+            return "keyword"
+
+        for name, register_cls in REGISTER_CLASSES.items():
+            with self.subTest(register=name):
+                registry = register_cls({"initial": initial_target}, keyword=keyword_target)
+
+                self.assertEqual(len(registry), 2)
+                self.assertEqual(list(registry), ["initial", "keyword"])
+                self.assertEqual(dict(registry), {"initial": initial_target, "keyword": keyword_target})
+                self.assertEqual(str(registry), str({"initial": initial_target, "keyword": keyword_target}))
+
+    def test_registers_named_and_unnamed_targets(self):
+        for name, register_cls in REGISTER_CLASSES.items():
+            with self.subTest(register=name):
+                registry = register_cls()
+
+                @registry("named")
+                def named_target():
+                    return "named"
+
+                @registry
+                def automatic_target():
+                    return "automatic"
+
+                self.assertIs(registry["named"], named_target)
+                self.assertIs(registry.get("automatic_target"), automatic_target)
+                self.assertEqual(set(registry.keys()), {"named", "automatic_target"})
+                self.assertEqual(set(registry.values()), {named_target, automatic_target})
+                self.assertEqual(set(registry.items()), {("named", named_target), ("automatic_target", automatic_target)})
+
+    def test_rejects_invalid_and_duplicate_targets(self):
+        for name, register_cls in REGISTER_CLASSES.items():
+            with self.subTest(register=name):
+                registry = register_cls()
+
+                def original_target():
+                    return "original"
+
+                with self.assertRaisesRegex(Exception, "must be callable"):
+                    registry.register("not callable")
+
+                registry("target")(original_target)
+                with self.assertRaisesRegex(Exception, "already exists"):
+                    registry("target")(lambda: None)
+                self.assertIs(registry["target"], original_target)
+
+    def test_supports_mutable_mapping_operations(self):
+        for name, register_cls in REGISTER_CLASSES.items():
+            with self.subTest(register=name):
+                registry = register_cls()
+
+                registry.update({"first": first_target})
+                registry.setdefault("second", second_target)
+                self.assertEqual(len(registry), 2)
+                self.assertIn("first", registry)
+
+                self.assertIs(registry.pop("first"), first_target)
+                del registry["second"]
+                self.assertFalse(registry)
+
+                registry["first"] = first_target
+                registry["second"] = second_target
+                registry.clear()
+                self.assertEqual(dict(registry), {})
+
+    def test_merges_registries_and_rejects_conflicts(self):
+        for name, register_cls in REGISTER_CLASSES.items():
+            with self.subTest(register=name):
+                registry = register_cls({"first": first_target})
+                other = register_cls({"second": second_target})
+
+                registry.merge(other)
+                self.assertEqual(dict(registry), {"first": first_target, "second": second_target})
+
+                with self.assertRaisesRegex(Exception, "already exists"):
+                    registry.merge(register_cls({"first": first_target}))
+                self.assertIs(registry["first"], first_target)
+
+
+if __name__ == "__main__":
+    unittest.main()