Source code for dantro.containers._registry

"""Implements a registry for dantro container types based on

import logging
from typing import Any, Dict, Optional, Union

from .._registry import ObjectRegistry as _ObjectRegistry
from ..base import BaseDataContainer
from ..exceptions import *

log = logging.getLogger(__name__)

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

[docs]class ContainerRegistry(_ObjectRegistry): _DESC = "container" _SKIP = False _OVERWRITE = True _EXPECTED_TYPE = (type,)
[docs] def _check_object(self, obj: Any) -> None: """Checks whether the object is valid.""" super()._check_object(obj) if not issubclass(obj, BaseDataContainer): raise TypeError( f"The given {self.desc} type {obj} needs to be a subclass of " "the dantro BaseDataContainer but was not!" )
[docs] def _register_via_decorator( self, obj: type, name: Optional[str] = None, **kws ): """Performs the registration operations when the decorator is used to register an object.""" self.register(obj, name=obj.__name__) self.register(obj, name=f"{obj.__module__}.{obj.__name__}") if name is not None: self.register(obj, name=name, **kws)
# ----------------------------------------------------------------------------- CONTAINERS = ContainerRegistry() """The dantro data container registry object."""
[docs]def register_container( Cls: type, name: str, *, skip_existing: bool = False, overwrite_existing: bool = True, ) -> None: """Adds an entry to the shared container registry. Args: Cls (type): The class that is to be registered as a container. name (str): The name to use for registration. skip_existing (bool, optional): Whether to skip registration if the container name is already registered. This suppresses the ValueError raised on existing container name. overwrite_existing (bool, optional): Whether to overwrite a potentially already existing container of the same name. If set, this takes precedence over ``skip_existing``. """ return CONTAINERS.register( Cls, name=name, skip_existing=skip_existing, overwrite_existing=overwrite_existing, )
[docs]def is_container( arg: Union[str, type] = None, /, **kws, ): """Decorator for registering containers with the container type registry. As an alternative to :py:func:`.register_container`, this decorator can be used to register a container right where its defined: .. testcode:: from dantro.containers import BaseDataContainer, is_container # Container name deduced from class name @is_container class MyDataContainer(BaseDataContainer): # ... do stuff here ... pass # Custom container name @is_container("my_container") class MyDataContainer(BaseDataContainer): # ... do stuff here ... pass # Overwriting a registered container of the same name @is_container("my_container", overwrite_existing=True) class MyDataContainer(BaseDataContainer): # ... do stuff here ... pass .. testcode:: :hide: from dantro.containers._registry import CONTAINERS assert MyDataContainer in CONTAINERS assert "MyDataContainer" in CONTAINERS assert "my_container" in CONTAINERS assert "PathContainer" in CONTAINERS """ return CONTAINERS._decorator(arg, **kws)