|
1 | | -from typing import Any, Callable, Dict, Type, TypeVar |
| 1 | +from typing import Callable, TypeVar |
2 | 2 |
|
3 | 3 | T = TypeVar("T") |
4 | 4 |
|
5 | | -REGISTRIES = {} |
6 | 5 |
|
| 6 | +class Registry(dict[str, type[T]]): |
| 7 | + """Registry that maps a name to its corresponding type.""" |
7 | 8 |
|
8 | | -def setup(registry_name: str): |
9 | | - """Setup a registry. |
| 9 | + def __init__(self, base_type: type[T]): |
| 10 | + super().__init__() |
| 11 | + REGISTRIES[base_type] = self |
| 12 | + self._base_type = base_type |
10 | 13 |
|
11 | | - Args: |
12 | | - registry_name (str): Registry name for grouping classes. |
| 14 | + def register(self, name: str) -> Callable[[type[T]], type[T]]: |
| 15 | + """Register a type as the given name. |
13 | 16 |
|
14 | | - Returns: |
15 | | - Tuple of the two functions: |
16 | | - - register: Register a class as the given name. |
17 | | - - get_cls: Return the registered class of the given name. |
18 | | - """ |
19 | | - REGISTRY = {} |
20 | | - REGISTRIES[registry_name] = REGISTRY |
| 17 | + Args: |
| 18 | + name (str): The name of a type. |
21 | 19 |
|
22 | | - def register(name: str) -> Callable[[Type[T]], Type[T]]: |
23 | | - """Register a class as the given name. |
| 20 | + Returns: |
| 21 | + Callable[[type[T]], type[T]]: Register decorator function. |
24 | 22 |
|
25 | | - Args: |
26 | | - name (str): The name of a class. |
| 23 | + Raises: |
| 24 | + ValueError: The type is already registered. |
27 | 25 | """ |
28 | 26 |
|
29 | | - def _register(cls: Type[T]): |
30 | | - if name in REGISTRY: |
| 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: |
31 | 32 | raise ValueError( |
32 | | - f"{name} already registered as {REGISTRY[name].__name__}. ({cls.__name__})" |
| 33 | + f"{cls.__name__}: `{name}` already registered as `{registered.__name__}`." |
33 | 34 | ) |
34 | | - REGISTRY[name] = cls |
| 35 | + self[name] = cls |
35 | 36 | return cls |
36 | 37 |
|
37 | 38 | return _register |
38 | 39 |
|
39 | | - def get_cls(name: str): |
40 | | - if name not in REGISTRY: |
41 | | - raise NotImplementedError( |
42 | | - f"`{name}` is not registered in `{registry_name}`." |
43 | | - ) |
44 | | - return REGISTRY[name] |
| 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 | + |
| 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 | + |
45 | 63 |
|
46 | | - return register, get_cls |
| 64 | +REGISTRIES: dict[type, Registry] = {} |
47 | 65 |
|
48 | 66 |
|
49 | | -def get_registry(registry_name: str) -> Dict[str, Type[Any]]: |
50 | | - """Get registry of the given name. |
| 67 | +def get_registry(base_type: type[T]) -> Registry[T]: |
| 68 | + """Get registry of the given base class type. |
51 | 69 |
|
52 | 70 | Args: |
53 | | - registry_name (str): Registry name to be returned. |
| 71 | + base_type (type[T]): Base class type that associated with the registry to be returned. |
54 | 72 |
|
55 | 73 | Returns: |
56 | | - Dict[str, Type[Any]]: Class mapper from registered name to its corresponding class. |
| 74 | + Registry[T]: Class mapper from registered name to its corresponding class. |
57 | 75 | """ |
58 | | - return REGISTRIES[registry_name] |
| 76 | + return REGISTRIES[base_type] |
0 commit comments