Source code for dantro._registry

"""Implements an object registry that can be specialized for certain use
cases, e.g. to store all available container types."""
import logging
from typing import Any, Optional, Union

from .exceptions import *
from .tools import make_columns
from .utils import KeyOrderedDict

log = logging.getLogger(__name__)

# -----------------------------------------------------------------------------


[docs]class ObjectRegistry: """""" _DESC: str = "object" """A description string for the entries of this registry""" _SKIP: bool = False """Default behavior for ``skip_existing`` argument""" _OVERWRITE: bool = False """Default behavior for ``overwrite_existing`` argument""" _EXPECTED_TYPE: Optional[Union[tuple, type]] = None """If set, will check for expected types""" # ......................................................................... def __init__(self): self._d = KeyOrderedDict() @property def classname(self) -> str: return self.__class__.__name__ @property def desc(self) -> str: return self._DESC def __getitem__(self, key: str) -> Any: try: return self._d[key] except KeyError: raise MissingRegistryEntry( f"Missing {self.desc} named '{key}'! " f"Available entries:\n{make_columns(self.keys())}" )
[docs] def keys(self): return self._d.keys()
[docs] def items(self): return self._d.items()
[docs] def values(self): return self._d.values()
def __iter__(self): return self._d.__iter__()
[docs] def __contains__(self, obj_or_key: Union[Any, str]) -> bool: """Whether the given argument is part of the keys or values of this registry.""" return obj_or_key in self._d.keys() or obj_or_key in self._d.values()
def __len__(self) -> int: return len(self._d) # .........................................................................
[docs] def _determine_name(self, obj: Any, *, name: Optional[str]) -> str: """Determines the object name, using a potentially given name""" if name is not None: return name if hasattr(obj, "__name__"): return obj.__name__ raise MissingNameError( f"Need `name` argument to register {self.desc} {obj} in " f"{self.classname}." )
[docs] def _check_object(self, obj: Any) -> None: """Checks whether the object is valid. If not, raises :py:exc:`~dantro.exceptions.InvalidRegistryEntry`. """ if self._EXPECTED_TYPE is not None: if not isinstance(obj, self._EXPECTED_TYPE): raise InvalidRegistryEntry( f"Expected type for {self.desc} in {self.classname} is " f"{self._EXPECTED_TYPE}, but got {type(obj)} " f"with value: '{obj}'" )
[docs] def register( self, obj: Any, name: Optional[str] = None, *, skip_existing: bool = None, overwrite_existing: bool = None, ) -> str: """Adds an entry to the registry. Args: obj (Any): The object to add to the registry. name (Optional[str], optional): The name to use. If not given, will deduce a name from the given object. skip_existing (bool, optional): Whether to skip registration if an object of that name already exists. If None, the classes default behavior (see :py:attr:`._SKIP`) is used. overwrite_existing (bool, optional): Whether to overwrite an entry if an object with that name already exists. If None, the classes default behavior (see :py:attr:`._OVERWRITE`) is used. """ if skip_existing is None: skip_existing = self._SKIP if overwrite_existing is None: overwrite_existing = self._OVERWRITE self._check_object(obj) name = self._determine_name(obj, name=name) if not isinstance(name, str): raise TypeError( f"{self.classname} `name` argument needs to be a string, but " f"was {type(name).__name__} with value: {name}" ) if name in self and not overwrite_existing: if skip_existing: log.debug( "A %s named '%s' is already registered and will not be " "registered again. Choose a different name or unset the " "`skip_existing` flag.", self.desc, name, ) return elif self[name] is obj: log.debug( "A %s named '%s' is already registered, but is identical " "to the existing one. Not setting it anew." ) return else: raise RegistryEntryExists( f"A {self.desc} named '{name}' is already registered in " f"{self.classname}! Not overwriting.\n" f" Existing: {self[name]}\n" f" New: {obj}\n" "Set the `overwrite_existing` flag to force overwriting " "or choose a different `name`. If no name was given, " "consider specifying it explicitly instead of letting it " "be deduced." ) self._d[name] = obj log.debug("Added %s '%s' to registry:\n%s", self.desc, name, obj) return name
# .. Decorator ............................................................
[docs] def _register_via_decorator(self, obj, name: Optional[str] = None, **kws): """Performs the registration operations when the decorator is used to register an object.""" self.register(obj, name=name, **kws)
[docs] def _decorator( self, arg: Union[Any, str] = None, /, **kws, ): """Method that can be used as a decorator for registering objects with this registry. Args: arg (Union[Any, str], optional): The name that should be used or the object that is to be added. If not a string, this refers to the ``@is_container`` call syntax **kws: Passed to :py:meth:`.register` """ def wrapper(obj): # Need some nonlocal names, see PEP3104 nonlocal arg if isinstance(arg, str): # Invocation via @is_operation() -- without arguments # The name was given as `arg` argument to _decorator kws["name"] = arg else: # Invocation via @is_operation -- without parentheses # No `arg` argument, thus no name information pass self._register_via_decorator(obj, **kws) return obj # Allow both invocation styles if arg is not None and not isinstance(arg, str): # @is_operation return wrapper(arg) # @is_operation(*args, **kwargs) return wrapper