File size: 1,934 Bytes
d661b19 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
# from detectron2
from typing import Any, Dict, Iterable, Iterator, Tuple
from tabulate import tabulate
class Registry(Iterable[Tuple[str, Any]]):
def __init__(self, name: str) -> None:
"""
Args:
name (str): the name of this registry
"""
self._name: str = name
self._obj_map: Dict[str, Any] = {}
def _do_register(self, name: str, obj: Any) -> None:
assert (
name not in self._obj_map
), "An object named '{}' was already registered in '{}' registry!".format(
name, self._name
)
self._obj_map[name] = obj
def register(self, obj: Any = None) -> Any:
"""
Register the given object under the the name `obj.__name__`.
Can be used as either a decorator or not. See docstring of this class for usage.
"""
if obj is None:
# used as a decorator
def deco(func_or_class: Any) -> Any:
name = func_or_class.__name__
self._do_register(name, func_or_class)
return func_or_class
return deco
# used as a function call
name = obj.__name__
self._do_register(name, obj)
def get(self, name: str) -> Any:
ret = self._obj_map.get(name)
if ret is None:
raise KeyError(
"No object named '{}' found in '{}' registry!".format(name, self._name)
)
return ret
def __contains__(self, name: str) -> bool:
return name in self._obj_map
def __repr__(self) -> str:
table_headers = ["Names", "Objects"]
table = tabulate(
self._obj_map.items(), headers=table_headers, tablefmt="fancy_grid"
)
return "Registry of {}:\n".format(self._name) + table
def __iter__(self) -> Iterator[Tuple[str, Any]]:
return iter(self._obj_map.items())
__str__ = __repr__
|