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,122 @@
+import sys
+import unittest
+from collections.abc import MutableMapping
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT))
+sys.path.insert(0, str(REPO_ROOT / "lightx2v_train"))
+
+import lightx2v.utils.registry_factory as runtime_registry
+import lightx2v_platform.registry_factory as platform_registry
+import lightx2v_train.utils.registry as training_registry
+
+
+def first_target():
+    return "first"
+
+
+def second_target():
+    return "second"
+
+
+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()