MuseV/musev/utils/register.py

import logging

logger = logging.getLogger(__name__)


class Register:
 def __init__(self, registry_name):
 self._dict = {}
 self._name = registry_name

 def __setitem__(self, key, value):
 if not callable(value):
 raise Exception(f"Value of a Registry must be a callable!\nValue: {value}")
 # 优先使用自定义的name，其次使用类名或者函数名。
 if "name" in value.__dict__:
 key = value.name
 elif key is None:
 key = value.__name__
 if key in self._dict:
 logger.warning("Key %s already in registry %s." % (key, self._name))
 self._dict[key] = value

 def register(self, target):
 """Decorator to register a function or class."""

 def add(key, value):
 self[key] = value
 return value

 if callable(target):
 # @reg.register
 return add(None, target)
 # @reg.register('alias')
 return lambda x: add(target, x)

 def __getitem__(self, key):
 return self._dict[key]

 def __contains__(self, key):
 return key in self._dict

 def keys(self):
 """key"""
 return self._dict.keys()