Source code for elasticai.creator.ir.registry

from collections.abc import Iterable, Iterator, Mapping
from typing import Self, overload

from .datagraph import DataGraph

_registry_types = []



[docs]
def is_registry(object) -> bool:
    return object in _registry_types




[docs]
def mark_as_registry[T](object: T) -> T:
    _registry_types.append(object)
    return object




[docs]
@mark_as_registry
class Registry[G: DataGraph](Mapping[str, G]):
    @overload
    def __init__(self, items: Iterable[tuple[str, G]], /) -> None: ...
    @overload
    def __init__(self, items: Mapping[str, G], /) -> None: ...
    @overload
    def __init__(self, /, **kwargs: G) -> None: ...

    def __init__(
        self,
        items: Iterable[tuple[str, G]] | Mapping[str, G] | None = None,
        /,
        **kwargs,
    ) -> None:
        if items is None:
            self._data: dict[str, G] = dict(**kwargs)
        else:
            if not len(kwargs) == 0:
                raise TypeError(f"unsupported arguments for {type(self)}")
            self._data = dict(items)


[docs]
    def __getitem__(self, name: str) -> G:
        return self._data[name]



[docs]
    def __len__(self) -> int:
        return len(self._data)



[docs]
    def __contains__(self, key: object) -> bool:
        return key in self._data



[docs]
    def __iter__(self) -> Iterator[str]:
        yield from self._data



[docs]
    def __or__(self, other: object) -> Self:
        if not isinstance(other, Mapping):
            raise TypeError("unsupported operand")
        return type(self)(**(self._data | dict(other.items())))  # type: ignore



[docs]
    def add(self, name: str, graph: G) -> Self:
        return type(self)(**(self._data) | {name: graph})