codekingpro/portable-devtools

1"""A registry of :class:`Schema <marshmallow.Schema>` classes. This allows for string2lookup of schemas, which may be used with3class:`fields.Nested <marshmallow.fields.Nested>`.4 5.. warning::6 7    This module is treated as private API.8    Users should not need to use this module directly.9"""10# ruff: noqa: ERA00111 12from __future__ import annotations13 14import typing15 16from marshmallow.exceptions import RegistryError17 18if typing.TYPE_CHECKING:19    from marshmallow import Schema20 21    SchemaType = type[Schema]22 23# {24#   <class_name>: <list of class objects>25#   <module_path_to_class>: <list of class objects>26# }27_registry = {}  # type: dict[str, list[SchemaType]]28 29 30def register(classname: str, cls: SchemaType) -> None:31    """Add a class to the registry of serializer classes. When a class is32    registered, an entry for both its classname and its full, module-qualified33    path are added to the registry.34 35    Example: ::36 37        class MyClass:38            pass39 40 41        register("MyClass", MyClass)42        # Registry:43        # {44        #   'MyClass': [path.to.MyClass],45        #   'path.to.MyClass': [path.to.MyClass],46        # }47 48    """49    # Module where the class is located50    module = cls.__module__51    # Full module path to the class52    # e.g. user.schemas.UserSchema53    fullpath = f"{module}.{classname}"54    # If the class is already registered; need to check if the entries are55    # in the same module as cls to avoid having multiple instances of the same56    # class in the registry57    if classname in _registry and not any(58        each.__module__ == module for each in _registry[classname]59    ):60        _registry[classname].append(cls)61    elif classname not in _registry:62        _registry[classname] = [cls]63 64    # Also register the full path65    if fullpath not in _registry:66        _registry.setdefault(fullpath, []).append(cls)67    else:68        # If fullpath does exist, replace existing entry69        _registry[fullpath] = [cls]70 71 72@typing.overload73def get_class(classname: str, *, all: typing.Literal[False] = ...) -> SchemaType: ...74 75 76@typing.overload77def get_class(78    classname: str, *, all: typing.Literal[True] = ...79) -> list[SchemaType]: ...80 81 82def get_class(classname: str, *, all: bool = False) -> list[SchemaType] | SchemaType:  # noqa: A00283    """Retrieve a class from the registry.84 85    :raises: `marshmallow.exceptions.RegistryError` if the class cannot be found86        or if there are multiple entries for the given class name.87    """88    try:89        classes = _registry[classname]90    except KeyError as error:91        raise RegistryError(92            f"Class with name {classname!r} was not found. You may need "93            "to import the class."94        ) from error95    if len(classes) > 1:96        if all:97            return _registry[classname]98        raise RegistryError(99            f"Multiple classes with name {classname!r} "100            "were found. Please use the full, "101            "module-qualified path."102        )103    return _registry[classname][0]104