thds.attrs-utils 1.6.20251103154246__py3-none-any.whl

thds/attrs_utils/jsonschema/util.py

@@ -0,0 +1,94 @@
+import functools
+import typing as ty
+from collections import ChainMap, defaultdict
+
+from ..type_cache import TypeCache
+from .constructors import INTEGER, NUMBER, TYPE
+
+ENUM, DEFAULT, ANYOF = (
+    "enum",
+    "default",
+    "anyOf",
+)
+DEFS, REF = "$defs", "$ref"
+TITLE = "title"
+
+JSONSchema = ty.Dict[str, ty.Any]
+JSON = ty.Union[int, float, bool, str, None, ty.List["JSON"], ty.Dict[str, "JSON"]]
+ToJSON = ty.Callable[[ty.Any], JSON]
+
+REF_TEMPLATE = f"#/{DEFS}/{{}}"
+
+
+class JSONSchemaTypeCache(TypeCache[JSONSchema]):
+    def to_defs(self):
+        return {self.type_names[type_id]: schema for type_id, schema in self.schemas.items()}
+
+    def to_ref(self, type_: ty.Type):
+        return {REF: REF_TEMPLATE.format(self.name_of(type_))}
+
+
+def check_cache(func):
+    """Decorator to cause a jsonschema generator to cache full schema definitions and return jsonschema
+    refs to the generated schemas. Useful for complex/custom types with names so as not to duplicate
+    sub-schemas throughout the main schema"""
+
+    @functools.wraps(func)
+    def new_func(gen_jsonschema, type_: ty.Type, type_cache: JSONSchemaTypeCache, serializer: ToJSON):
+        if type_ not in type_cache:
+            type_name = type_cache.name_of(type_)
+            schema = {TITLE: type_name, **func(gen_jsonschema, type_, type_cache, serializer)}
+            type_cache[type_] = schema
+
+        return type_cache.to_ref(type_)
+
+    return new_func
+
+
+def _flatten_anyof_schemas(schemas: ty.Iterable[JSONSchema]) -> ty.Iterator[JSONSchema]:
+    for schema in schemas:
+        if ANYOF in schema and len(schema) == 1:
+            yield from _flatten_anyof_schemas(schema[ANYOF])
+        else:
+            yield schema
+
+
+def _merge_schemas_anyof(schemas: ty.Iterable[JSONSchema]) -> JSONSchema:
+    enum_values = []
+    schemas_by_type = defaultdict(list)
+    other_schemas = []
+    number = "_"
+    for schema in _flatten_anyof_schemas(schemas):
+        if REF in schema:
+            other_schemas.append(schema)
+        elif ENUM in schema and len(schema) == 1:
+            # enums can simply be combined
+            enum_values.extend(schema[ENUM])
+        elif TYPE in schema:
+            types = [schema[TYPE]] if isinstance(schema[TYPE], str) else schema[TYPE]
+            for t in types:
+                # integer and number share keywords; combine into the same list
+                if t in (INTEGER, NUMBER):
+                    t = number
+                schemas_by_type[t].append(schema)
+        else:
+            other_schemas.append(schema)
+
+    # any type with only one schema can be keyword-merged with all other such schemas
+    simple = []
+    complex = []
+    for _type, schemas in schemas_by_type.items():
+        if len(schemas) == 1:
+            simple.append(schemas[0].copy())
+        else:
+            # we'll transfer these unchanged
+            complex.extend(schemas)
+
+    simple_types = [schema.pop(TYPE) for schema in simple]
+    simple_type = simple_types[0] if len(simple_types) == 1 else simple_types
+    multi_type_schema = {TYPE: simple_type, **ChainMap(*simple)} if simple else None
+    enum_schema = {ENUM: enum_values} if enum_values else None
+    all_schemas = [schema for schema in (enum_schema, multi_type_schema) if schema is not None]
+    all_schemas.extend(complex)
+    all_schemas.extend(other_schemas)
+    return all_schemas[0] if len(all_schemas) == 1 else {ANYOF: all_schemas}

thds/attrs_utils/random/__init__.py

@@ -0,0 +1,3 @@
+# ensure registrations
+from . import attrs, builtin, collection, optional, registry, tuple, union, util  # noqa: F401
+from .gen import random_gen  # noqa: F401

thds/attrs_utils/random/attrs.py

@@ -0,0 +1,57 @@
+from typing import Callable, Dict, Sequence, Type, TypeVar
+
+import attrs
+
+from ..registry import Registry
+from ..type_utils import is_namedtuple_type
+from .util import Gen, T
+
+AT = TypeVar("AT", bound=attrs.AttrsInstance)
+
+
+class GenRecordByFieldNameRegistry(Registry[Type, Dict[str, Gen]]):
+    # this class def only exists because of the mypy error "Type variable is unbound"
+    pass
+
+
+CUSTOM_ATTRS_BY_FIELD_REGISTRY = GenRecordByFieldNameRegistry()
+
+
+def random_attrs(
+    constructor: Callable[..., AT], arg_gens: Sequence[Gen[T]], kwarg_gens: Dict[str, Gen[T]]
+) -> AT:
+    return constructor(*(gen() for gen in arg_gens), **{name: gen() for name, gen in kwarg_gens.items()})
+
+
+def _register_random_gen_by_field(type_: Type[T], **gens: Gen):
+    if attrs.has(type_):
+        fields = attrs.fields(type_)  # type: ignore
+        names = [f.name for f in fields]
+    elif is_namedtuple_type(type_):
+        names = type_._fields  # type: ignore [attr-defined]
+    else:
+        raise TypeError(f"Don't know how to interpret {type_} as a record type")
+
+    unknown_names = set(gens).difference(names)
+    assert not unknown_names, f"Unknown fields: {unknown_names}"
+
+    CUSTOM_ATTRS_BY_FIELD_REGISTRY.register(type_, gens)
+
+
+def register_random_gen_by_field(**gens: Gen):
+    """Register random generators for the fields of a record type by specifying random generators by name
+
+    Example:
+
+        from typing import NamedTuple
+
+        @register_random_gen_by_field(foo=lambda: 42):
+        class Bar(NamedTuple):
+            foo: int
+    """
+
+    def decorator(type_: Type[AT]) -> Gen[AT]:
+        _register_random_gen_by_field(type_, **gens)
+        return type_
+
+    return decorator

thds/attrs_utils/random/builtin.py

@@ -0,0 +1,103 @@
+import datetime
+import random
+import string
+from functools import partial
+
+from .util import Gen
+
+MIN_INT = -(2**32)
+MAX_INT = 2**32 - 1
+MIN_FLOAT = -float(MIN_INT)
+MAX_FLOAT = float(MAX_INT)
+MIN_DATE = datetime.date(1970, 1, 1)
+MAX_DATE = datetime.date.today()
+STR_MAX_LEN = 32
+
+
+def random_null():
+    return None
+
+
+def random_int(lo: int = MIN_INT, hi: int = MAX_INT) -> int:
+    return random.randint(lo, hi)
+
+
+def random_int_gen(lo: int = MIN_INT, hi: int = MAX_INT) -> Gen[int]:
+    return partial(random.randint, lo, hi)
+
+
+def random_float(lo: float = MIN_FLOAT, hi: float = MAX_FLOAT) -> float:
+    return random.uniform(lo, hi)
+
+
+def random_float_gen(lo: float = MIN_FLOAT, hi: float = MAX_FLOAT) -> Gen[float]:
+    return partial(random.uniform, lo, hi)
+
+
+default_str_len_gen = random_int_gen(0, STR_MAX_LEN)
+
+
+def random_str(len_gen: Gen[int] = default_str_len_gen, chars: str = string.printable) -> str:
+    return "".join(random.choices(chars, k=len_gen()))
+
+
+def random_str_gen(len_gen: Gen[int] = default_str_len_gen, chars: str = string.printable) -> Gen[str]:
+    return partial(random_str, len_gen, chars)
+
+
+def random_bool(true_rate: float = 0.5) -> bool:
+    return random.random() < true_rate
+
+
+def random_bool_gen(true_rate: float = 0.5) -> Gen[bool]:
+    return partial(random_bool, true_rate)
+
+
+default_date_offset_gen = random_int_gen(0, (MAX_DATE - MIN_DATE).days)
+
+
+def random_date(
+    earliest: datetime.date = MIN_DATE,
+    offset: Gen[int] = default_date_offset_gen,
+) -> datetime.date:
+    return earliest + datetime.timedelta(days=offset())
+
+
+def random_date_gen(
+    earliest: datetime.date,
+    offset: Gen[int] = default_date_offset_gen,
+) -> Gen[datetime.date]:
+    return partial(random_date, earliest, offset)
+
+
+def random_date_gen_from_range(earliest: datetime.date, latest: datetime.date):
+    return random_date_gen(earliest, random_int_gen(0, (latest - earliest).days))
+
+
+default_datetime_offset_gen = random_float_gen(0.0, (MAX_DATE - MIN_DATE).total_seconds())
+
+
+def random_datetime(
+    earliest: datetime.date = MIN_DATE,
+    offset: Gen[float] = default_datetime_offset_gen,
+) -> datetime.date:
+    return earliest + datetime.timedelta(seconds=offset())
+
+
+def random_datetime_gen(
+    earliest: datetime.datetime,
+    offset: Gen[float] = default_datetime_offset_gen,
+) -> Gen[datetime.date]:
+    return partial(random_datetime, earliest, offset)
+
+
+def random_datetime_gen_from_range(earliest: datetime.datetime, latest: datetime.datetime):
+    return random_datetime_gen(earliest, random_float_gen(0.0, (latest - earliest).total_seconds()))
+
+
+def random_bytes(len_gen: Gen[int] = default_str_len_gen):
+    return bytes(random.randint(0, 255) for _ in range(len_gen()))
+
+
+def random_bytearray(len_gen: Gen[int] = default_str_len_gen):
+    return bytearray(random.randint(0, 255) for _ in range(len_gen()))

thds/attrs_utils/random/collection.py

@@ -0,0 +1,41 @@
+import itertools
+from functools import partial
+from typing import Callable, Collection, Iterable, Mapping, Tuple, TypeVar, cast
+
+from .builtin import random_int_gen
+from .util import Gen
+
+T = TypeVar("T")
+U = TypeVar("U")
+C = TypeVar("C", bound=Collection)
+M = TypeVar("M", bound=Mapping)
+KV = Tuple[T, U]
+
+COLLECTION_MIN_LEN = 0
+COLLECTION_MAX_LEN = 32
+
+default_len_gen = random_int_gen(COLLECTION_MIN_LEN, COLLECTION_MAX_LEN)
+
+
+def random_collection(
+    constructor: Callable[[Iterable[T]], C],
+    value_gen: Gen[Iterable[T]],
+    len_gen: Gen[int] = default_len_gen,
+) -> C:
+    return constructor(itertools.islice(value_gen(), len_gen()))
+
+
+def random_collection_gen(
+    constructor: Callable[[Iterable[T]], C],
+    value_gen: Gen[Iterable[T]],
+    len_gen: Gen[int] = default_len_gen,
+) -> Gen[C]:
+    return partial(random_collection, constructor, value_gen, len_gen)
+
+
+def random_mapping_gen(
+    constructor: Callable[[Iterable[KV]], M],
+    kv_gen: Gen[Iterable[KV]],
+    len_gen: Gen[int] = default_len_gen,
+) -> Gen[M]:
+    return cast(Gen[M], partial(random_collection, constructor, kv_gen, len_gen))

thds/attrs_utils/random/gen.py

@@ -0,0 +1,134 @@
+import collections
+import enum
+from functools import partial
+from typing import DefaultDict, Iterable, Tuple, Type, cast, get_args
+
+import attr
+
+from .. import recursion, type_recursion, type_utils
+from . import attrs, collection, optional, tuple, union
+from .registry import GEN_REGISTRY
+from .util import Gen, T, U, choice_gen, juxtapose_gen, repeat_gen
+
+_UNKNOWN_TYPE = (
+    "Don't know how to generate random instances of {!r}; use "
+    f"{__name__}.random_gen.register to register a random generator"
+)
+
+unknown_type: "recursion.RecF[Type, [], Gen]" = recursion.value_error(_UNKNOWN_TYPE, TypeError)
+
+
+def gen_literal(random_gen, type_: Type[T]) -> Gen[T]:
+    values = cast(Tuple[T], get_args(type_))
+    return choice_gen(values)
+
+
+def gen_enum(random_gen, type_: Type[T]) -> Gen[T]:
+    assert issubclass(type_, enum.Enum)
+    return choice_gen(list(type_))
+
+
+def gen_attrs(random_gen, type_: Type[attr.AttrsInstance]) -> Gen[attr.AttrsInstance]:
+    fields = attr.fields(type_)  # type: ignore [arg-type,misc]
+    kw_only_fields = [f for f in fields if f.kw_only]
+    overrides = attrs.CUSTOM_ATTRS_BY_FIELD_REGISTRY.get(type_)
+
+    def random_gen_(field: attr.Attribute):
+        if overrides:
+            return overrides.get(field.name) or random_gen(field.type)
+        return random_gen(field.type)
+
+    if kw_only_fields:
+        pos_fields = [f for f in fields if not f.kw_only]
+        arg_gens = list(map(random_gen_, pos_fields))
+        kwarg_gens = {f.name: random_gen_(f) for f in kw_only_fields}
+        return partial(attrs.random_attrs, type_, arg_gens, kwarg_gens)
+    else:
+        field_gens = map(random_gen_, fields)
+        return tuple.random_namedtuple_gen(type_, *field_gens)
+
+
+def gen_namedtuple(random_gen, type_: Type[T]) -> Gen[T]:
+    field_names = type_._fields  # type: ignore [attr-defined]
+    field_types = (type_.__annotations__[name] for name in field_names)  # type: ignore [attr-defined]
+    overrides = attrs.CUSTOM_ATTRS_BY_FIELD_REGISTRY.get(type_)
+    if overrides:
+        return tuple.random_namedtuple_gen(
+            type_, *(overrides.get(name) or random_gen(t) for name, t in zip(field_types, field_names))
+        )
+    else:
+        return tuple.random_namedtuple_gen(type_, *map(random_gen, field_types))
+
+
+def gen_tuple(random_gen, type_: Type[T]) -> Gen[T]:
+    args = get_args(type_)
+    if not args:
+        raise TypeError(_UNKNOWN_TYPE.format(type_))
+    return cast(Gen[T], tuple.random_tuple_gen(*map(random_gen, args)))
+
+
+def gen_variadic_tuple(random_gen, type_: Type[T]) -> Gen[T]:
+    args = get_args(type_)
+    cons = type_utils.concrete_constructor(type_)
+    v_gen = random_gen(args[0])
+    return cast(Gen[T], collection.random_collection_gen(cons, repeat_gen(v_gen)))
+
+
+def gen_optional(random_gen, type_: Type[T]) -> Gen[T]:
+    # more specialized that the union case just below
+    return optional.random_optional_gen(
+        random_gen(type_utils.unwrap_optional(type_))
+    )  # type: ignore [return-value]
+
+
+def gen_union(random_gen, type_: Type[T]) -> Gen[T]:
+    return union.random_uniform_union_gen(*map(random_gen, get_args(type_)))
+
+
+def _construct_defauldict(d: DefaultDict[T, U], kvs: Iterable[Tuple[T, U]]) -> DefaultDict[T, U]:
+    for k, v in kvs:
+        d[k] = v
+    return d
+
+
+def gen_mapping(random_gen, type_: Type[collection.M]) -> Gen[collection.M]:
+    # more specific that collection case below
+    args = get_args(type_)
+    if len(args) != 2:
+        raise TypeError(_UNKNOWN_TYPE.format(type_))
+    kt, vt = args
+    v_gen = random_gen(vt)
+    kv_gen = juxtapose_gen(random_gen(kt), v_gen)
+
+    cons = (
+        partial(_construct_defauldict, collections.defaultdict(v_gen))
+        if type_utils.get_origin(type_) is type_utils.get_origin(DefaultDict)
+        else type_utils.concrete_constructor(type_)  # type: ignore [arg-type]
+    )
+    return collection.random_mapping_gen(cons, repeat_gen(kv_gen))  # type: ignore [arg-type]
+
+
+def gen_collection(random_gen, type_: Type[collection.C]) -> Gen[collection.C]:
+    # most generic collection case
+    args = get_args(type_)
+    if len(args) != 1:
+        raise TypeError(_UNKNOWN_TYPE.format(type_))
+    cons = type_utils.concrete_constructor(type_)
+    v_gen = random_gen(args[0])
+    return cast(Gen[collection.C], collection.random_collection_gen(cons, repeat_gen(v_gen)))
+
+
+random_gen: "type_recursion.ConstructorFactory[[]]" = type_recursion.ConstructorFactory(
+    GEN_REGISTRY,
+    attrs=gen_attrs,
+    namedtuple=gen_namedtuple,
+    literal=gen_literal,
+    enum=gen_enum,
+    optional=gen_optional,
+    union=gen_union,
+    tuple=gen_tuple,
+    variadic_tuple=gen_variadic_tuple,
+    mapping=gen_mapping,
+    collection=gen_collection,
+    otherwise=unknown_type,
+)

thds/attrs_utils/random/optional.py

@@ -0,0 +1,13 @@
+from functools import partial
+from typing import Optional
+
+from .builtin import random_bool
+from .util import Gen, T
+
+
+def random_optional(gen: Gen[T], nonnull: Gen[bool] = random_bool) -> Optional[T]:
+    return gen() if nonnull() else None
+
+
+def random_optional_gen(gen: Gen[T], nonnull: Gen[bool] = random_bool) -> Gen[Optional[T]]:
+    return partial(random_optional, gen, nonnull)

thds/attrs_utils/random/registry.py

@@ -0,0 +1,30 @@
+import datetime
+from typing import Type
+
+from ..registry import Registry
+from . import builtin
+from .util import Gen, T
+
+
+class GenTypeRegistry(Registry[Type[T], Gen[T]]):
+    # This class def only exists because of the mypy error "Type variable is unbound".
+    # We use the typevar here however to assert that the registered generators generate instances of the
+    # types they're registered to
+    pass
+
+
+# we define the registry here before passing it to the recursion so that the recursion can reference it
+# after it's already registered implementations in the various type-specific sibling modules to this one
+GEN_REGISTRY: GenTypeRegistry = GenTypeRegistry(
+    [
+        (type(None), builtin.random_null),
+        (int, builtin.random_int),
+        (bool, builtin.random_bool),
+        (float, builtin.random_float),
+        (str, builtin.random_str),
+        (bytes, builtin.random_bytes),
+        (bytearray, builtin.random_bytearray),
+        (datetime.date, builtin.random_date),
+        (datetime.datetime, builtin.random_datetime),
+    ]
+)

thds/attrs_utils/random/tuple.py

@@ -0,0 +1,24 @@
+from functools import partial
+from typing import Callable, Sequence, Tuple, TypeVar
+
+from .util import Gen, T
+
+N = TypeVar("N")
+
+
+def random_tuple(gens: Sequence[Gen[T]]) -> Tuple[T, ...]:
+    # can't type this very well for heterogeneous tuples without going a little crazy.
+    # most of the time it's being generated from types in an automated way so it's not a worry.
+    return tuple(gen() for gen in gens)
+
+
+def random_tuple_gen(*gens: Gen[T]) -> Gen[Tuple[T, ...]]:
+    return partial(random_tuple, gens)
+
+
+def random_namedtuple(constructor: Callable[..., N], gens: Sequence[Gen[T]]) -> N:
+    return constructor(*random_tuple(gens))
+
+
+def random_namedtuple_gen(constructor: Callable[..., N], *gens: Gen[T]) -> Gen[N]:
+    return partial(random_namedtuple, constructor, gens)

thds/attrs_utils/random/union.py

@@ -0,0 +1,33 @@
+import random
+from functools import partial
+from itertools import accumulate
+from operator import add
+from typing import Sequence, Tuple
+
+from .util import Gen, T
+
+
+def random_weighted_union(gens: Sequence[Gen[T]], cum_weights: Sequence[float]):
+    """optimized to use the cum_weights arg of random.choices"""
+    gen = random.choices(gens, cum_weights=cum_weights, k=1)[0]
+    return gen()
+
+
+def random_weighted_union_gen(*weighted_gens: Tuple[Gen[T], float]) -> Gen[T]:
+    # can't type this very well for heterogeneous tuples without going a little crazy.
+    # most of the time it's being generated from types in an automated way so it's not a worry.
+    gens, weights = zip(*weighted_gens)
+    total = sum(weights)
+    cum_weights = tuple(w / total for w in accumulate(weights, add))
+    return partial(random_weighted_union, gens, cum_weights)
+
+
+def random_uniform_union(gens: Sequence[Gen[T]]) -> T:
+    gen = random.choice(gens)
+    return gen()
+
+
+def random_uniform_union_gen(*gens: Gen[T]) -> Gen[T]:
+    # can't type this very well for heterogeneous tuples without going a little crazy.
+    # most of the time it's being generated from types in an automated way so it's not a worry.
+    return partial(random_uniform_union, gens)

thds/attrs_utils/random/util.py

@@ -0,0 +1,55 @@
+import random
+from functools import partial
+from itertools import accumulate
+from operator import add
+from typing import Callable, Iterator, Mapping, Optional, Sequence, Tuple, TypeVar, Union
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+Gen = Callable[[], T]
+
+
+def repeat(gen: Gen[T]) -> Iterator[T]:
+    while True:
+        yield gen()
+
+
+def repeat_gen(gen: Gen[T]) -> Gen[Iterator[T]]:
+    return partial(repeat, gen)
+
+
+def juxtapose(gen1: Gen[T], gen2: Gen[U]) -> Tuple[T, U]:
+    # e.g. for key-value pairs
+    return gen1(), gen2()
+
+
+def juxtapose_gen(gen1: Gen[T], gen2: Gen[U]) -> Gen[Tuple[T, U]]:
+    return partial(juxtapose, gen1, gen2)
+
+
+def either(gen1: Gen[T], gen2: Gen[U], choose1: Gen[bool]) -> Union[T, U]:
+    return gen1() if choose1() else gen2()
+
+
+def either_gen(gen1: Gen[T], gen2: Gen[U], choose1: Gen[bool]) -> Gen[Union[T, U]]:
+    return partial(either, gen1, gen2, choose1)
+
+
+def choice(values: Sequence[T], cum_weights: Optional[Sequence[float]] = None) -> T:
+    """optimized to use the cum_weights arg of random.choices"""
+    return (
+        random.choice(values)
+        if cum_weights is None
+        else random.choices(values, cum_weights=cum_weights, k=1)[0]
+    )
+
+
+def choice_gen(values: Union[Sequence[T], Mapping[T, float]]) -> Gen[T]:
+    if isinstance(values, Mapping):
+        values_, weights = zip(*values.items())
+        total = sum(weights)
+        cum_weights = tuple(w / total for w in accumulate(weights, add))
+        return partial(choice, values_, cum_weights)
+    else:
+        return partial(choice, values)

thds/attrs_utils/recursion.py

@@ -0,0 +1,48 @@
+from functools import partial
+from typing import Any, Callable, Generic, Sequence, Tuple, Type, TypeVar
+
+from typing_extensions import Concatenate, ParamSpec
+
+T = TypeVar("T", contravariant=True)
+U = TypeVar("U", covariant=True)
+Params = ParamSpec("Params")
+
+Predicate = Callable[[T], bool]
+F = Callable[Concatenate[T, Params], U]
+RecF = Callable[Concatenate[F, T, Params], U]
+
+
+class StructuredRecursion(Generic[T, Params, U]):
+    def __init__(
+        self,
+        guarded_recursions: Sequence[Tuple[Predicate[T], RecF[T, Params, U]]],
+        fallback: RecF[T, Params, U],
+    ):
+        self.guarded_recursions = list(guarded_recursions)
+        self.fallback = fallback
+
+    def __call__(self, obj: T, *args: Params.args, **kwargs: Params.kwargs) -> U:
+        for predicate, recurse in self.guarded_recursions:
+            if predicate(obj):
+                return recurse(self, obj, *args, **kwargs)
+        else:
+            return self.fallback(self, obj, *args, **kwargs)
+
+
+def _value_error(
+    msg: str,
+    exc_type: Type[Exception],
+    # placeholder for the recursive function in case you wish to specify this explictly as one of the
+    # recursions; type doesn't matter
+    f: Any,
+    obj: T,
+    *args,
+    **kwargs,
+):
+    raise exc_type(msg.format(obj))
+
+
+def value_error(msg: str, exc_type: Type[Exception] = ValueError) -> RecF[T, Params, U]:
+    """Helper to be passed as the `fallback` of a `StructuredRecursion` in case there is no natural
+    fallback/default implementation and an input fails to satisfy any of the predicates"""
+    return partial(_value_error, msg, exc_type)

thds/attrs_utils/registry.py

@@ -0,0 +1,40 @@
+from functools import partial, wraps
+from typing import Callable, Dict, Type, TypeVar, cast, overload
+
+T = TypeVar("T")
+U = TypeVar("U")
+_D = TypeVar("_D")
+_Type = TypeVar("_Type", bound=Type)
+
+_MISSING = object()
+
+
+class Registry(Dict[T, U]):
+    @overload
+    def register(self, key: T) -> Callable[[U], U]: ...
+
+    @overload
+    def register(self, key: T, value: U) -> U: ...
+
+    def register(self, key: T, value=_MISSING):
+        if value is _MISSING:
+
+            def decorator(value: U) -> U:
+                self[key] = value
+                return value
+
+            return decorator
+        else:
+            self[key] = value
+            return value
+
+    def cache(self, func: Callable[[T], U]) -> Callable[[T], U]:
+        cached = partial(_check_cache, self, func)
+        return wraps(func)(cached)  # type: ignore [return-value]
+
+
+def _check_cache(cache: Dict[T, U], func: Callable[[T], U], key: T) -> U:
+    value = cache.get(key, cast(U, _MISSING))
+    if value is _MISSING:
+        value = cache[key] = func(key)
+    return value