diff --git a/lightx2v/utils/registry_factory.py b/lightx2v/utils/registry_factory.py index dabd48e80..69102417b 100755 --- a/lightx2v/utils/registry_factory.py +++ b/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 diff --git a/lightx2v_platform/base/intel_xpu.py b/lightx2v_platform/base/intel_xpu.py index cc6c574ac..db4a665e3 100644 --- a/lightx2v_platform/base/intel_xpu.py +++ b/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 diff --git a/lightx2v_platform/registry_factory.py b/lightx2v_platform/registry_factory.py index 06a9d5404..478570bd0 100755 --- a/lightx2v_platform/registry_factory.py +++ b/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 diff --git a/lightx2v_train/lightx2v_train/utils/registry.py b/lightx2v_train/lightx2v_train/utils/registry.py index ef129e83e..e50cc25bd 100644 --- a/lightx2v_train/lightx2v_train/utils/registry.py +++ b/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 diff --git a/test_cases/test_registry_factory.py b/test_cases/test_registry_factory.py new file mode 100644 index 000000000..5551927be --- /dev/null +++ b/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()