Source code for dantro.groups._registry

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

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

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

log = logging.getLogger(__name__)

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

[docs]class GroupRegistry(_ObjectRegistry): _DESC = "group" _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, BaseDataGroup): 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)
# ----------------------------------------------------------------------------- GROUPS = GroupRegistry() """The dantro data group registry object."""
[docs]def register_group( Cls: type, name: str, *, skip_existing: bool = False, overwrite_existing: bool = True, ) -> None: """Adds an entry to the shared group registry. Args: Cls (type): The class that is to be registered as a group. name (str): The name to use for registration. skip_existing (bool, optional): Whether to skip registration if the group name is already registered. This suppresses the ValueError raised on existing group name. overwrite_existing (bool, optional): Whether to overwrite a potentially already existing group of the same name. If set, this takes precedence over ``skip_existing``. """ return GROUPS.register( Cls, name=name, skip_existing=skip_existing, overwrite_existing=overwrite_existing, )
[docs]def is_group( arg: Union[str, type] = None, /, **kws, ): """Decorator for registering groups with the container type registry. As an alternative to :py:func:`.register_group`, this decorator can be used to register a container right where its defined: .. testcode:: from dantro.groups import BaseDataGroup, is_group # Group name deduced from class name @is_group class MyDataGroup(BaseDataGroup): # ... do stuff here ... pass # Custom group name @is_group("my_group") class MyDataGroup(BaseDataGroup): # ... do stuff here ... pass # Overwriting a registered group of the same name @is_group("my_group", overwrite_existing=True) class MyDataGroup(BaseDataGroup): # ... do stuff here ... pass .. testcode:: :hide: from dantro.groups._registry import GROUPS assert MyDataGroup in GROUPS assert "MyDataGroup" in GROUPS assert "my_group" in GROUPS assert "OrderedDataGroup" in GROUPS """ return GROUPS._decorator(arg, **kws)