-
Notifications
You must be signed in to change notification settings - Fork 1
/
registry.py
81 lines (68 loc) · 2.51 KB
/
registry.py
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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# pyre-ignore-all-errors[2,3]
from typing import Any, Dict, Iterable, Iterator, Tuple
from tabulate import tabulate
class Registry(Iterable[Tuple[str, Any]]):
"""
The registry that provides name -> object mapping, to support third-party
users' custom modules.
To create a registry (e.g. a backbone registry):
.. code-block:: python
BACKBONE_REGISTRY = Registry('BACKBONE')
To register an object:
.. code-block:: python
@BACKBONE_REGISTRY.register()
class MyBackbone():
...
Or:
.. code-block:: python
BACKBONE_REGISTRY.register(MyBackbone)
"""
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, prefix: str = "") -> 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(prefix + name, func_or_class)
return func_or_class
return deco
# used as a function call
name = obj.__name__
self._do_register(prefix + 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())
# pyre-fixme[4]: Attribute must be annotated.
__str__ = __repr__