Source code for mbrs.registry

 1from typing import Callable, TypeVar
 2
 3T = TypeVar("T")
 4
 5

[docs]
 6class Registry(dict[str, type[T]]):
 7    """Registry that maps a name to its corresponding type."""
 8
 9    def __init__(self, base_type: type[T]):
10        super().__init__()
11        REGISTRIES[base_type] = self
12        self._base_type = base_type
13

[docs]
14    def register(self, name: str) -> Callable[[type[T]], type[T]]:
15        """Register a type as the given name.
16
17        Args:
18            name (str): The name of a type.
19
20        Returns:
21            Callable[[type[T]], type[T]]: Register decorator function.
22
23        Raises:
24            ValueError: The type is already registered.
25        """
26
27        def _register(cls: type[T]) -> type[T]:
28            if not issubclass(cls, self._base_type):
29                raise ValueError(f"`{cls.__name__}` must inherit `{self._base_type}`.")
30
31            if (registered := self.get(name)) is not None:
32                raise ValueError(
33                    f"{cls.__name__}: `{name}` already registered as `{registered.__name__}`."
34                )
35            self[name] = cls
36            return cls
37
38        return _register

39

[docs]
40    def get_cls(self, name: str) -> type[T]:
41        """Get a class type.
42
43        Args:
44            name: A registered name.
45
46        Returns:
47            type[T]: Class type.
48        """
49        return self.__getitem__(name)

50

[docs]
51    def get_closure(
52        self,
53    ) -> tuple[Callable[[str], Callable[[type[T]], type[T]]], Callable[[str], type[T]]]:
54        """Get closure functions: `register()` and `get_cls()`.
55
56        Returns:
57            tuple:
58              - Callable[[str], Callable[[type[T]], type[T]]]: `register()` function.
59              - Callable[[str], type[T]]: `get_cls()` function.
60        """
61        return (self.register, self.get_cls)


62
63
64REGISTRIES: dict[type, Registry] = {}
65
66

[docs]
67def get_registry(base_type: type[T]) -> Registry[T]:
68    """Get registry of the given base class type.
69
70    Args:
71        base_type (type[T]): Base class type that associated with the registry to be returned.
72
73    Returns:
74        Registry[T]: Class mapper from registered name to its corresponding class.
75    """
76    return REGISTRIES[base_type]