diff --git a/glue_ar/registries.py b/glue_ar/registries.py index d7bdb3d..742a2bd 100644 --- a/glue_ar/registries.py +++ b/glue_ar/registries.py @@ -19,6 +19,7 @@ def add(self, extensions: Union[str, Tuple[str]], builder: Type): def __call__(self, extensions: Union[str, Tuple[str]]): def adder(builder: Type): self.add(extensions, builder) + return builder return adder @@ -33,6 +34,7 @@ def add(self, name: str, compressor: Callable[[str], None]): def __call__(self, name: str): def adder(compressor: Callable[[str], None]): self.add(name, compressor) + return compressor return adder diff --git a/glue_ar/tests/__init__.py b/glue_ar/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/glue_ar/tests/test_registries.py b/glue_ar/tests/test_registries.py new file mode 100644 index 0000000..ed2070d --- /dev/null +++ b/glue_ar/tests/test_registries.py @@ -0,0 +1,29 @@ +from ..registries import BuilderRegistry, CompressorRegistry + + +def test_builder_registry(): + + builder_registry = BuilderRegistry() + + @builder_registry("ext") + class ExtBuilder: + pass + + @builder_registry(("ext1", "ext2")) + class ExtsBuilder: + pass + + assert builder_registry.members == {"ext": ExtBuilder, + "ext1": ExtsBuilder, + "ext2": ExtsBuilder} + + +def test_compressor_registry(): + + compressor_registry = CompressorRegistry() + + @compressor_registry("dummy") + def dummy_compressor(_filepath: str): + pass + + assert compressor_registry.members == {"dummy": dummy_compressor}