From 1cb5a18a6cc8465795fafe98b49fb7f9d1b201bf Mon Sep 17 00:00:00 2001 From: GlowLED Date: Sun, 12 Jul 2026 00:23:39 +0800 Subject: [PATCH 1/2] refactor: implement registries as mutable mappings --- lightx2v/utils/registry_factory.py | 16 +- lightx2v_platform/base/intel_xpu.py | 2 +- lightx2v_platform/registry_factory.py | 17 +- .../lightx2v_train/utils/registry.py | 17 +- test_cases/test_registry_factory.py | 147 ++++++++++++++++++ 5 files changed, 189 insertions(+), 10 deletions(-) create mode 100644 test_cases/test_registry_factory.py 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..3aa12226a --- /dev/null +++ b/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() From b6e578607c2053b513f04c641b92fa7026c51b42 Mon Sep 17 00:00:00 2001 From: GlowLED Date: Sun, 12 Jul 2026 11:40:37 +0800 Subject: [PATCH 2/2] Update test_cases/test_registry_factory.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- test_cases/test_registry_factory.py | 35 +++++------------------------ 1 file changed, 5 insertions(+), 30 deletions(-) diff --git a/test_cases/test_registry_factory.py b/test_cases/test_registry_factory.py index 3aa12226a..5551927be 100644 --- a/test_cases/test_registry_factory.py +++ b/test_cases/test_registry_factory.py @@ -1,21 +1,15 @@ -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] +sys.path.insert(0, str(REPO_ROOT)) +sys.path.insert(0, str(REPO_ROOT / "lightx2v_train")) - -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 +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(): @@ -26,25 +20,6 @@ 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,