haiway 0.10.14__py3-none-any.whl → 0.10.16__py3-none-any.whl

haiway/state/attributes.py

@@ -0,0 +1,747 @@
+import types
+import typing
+from collections.abc import Callable, Mapping, MutableMapping, Sequence
+from types import GenericAlias, NoneType, UnionType
+from typing import (
+    Any,
+    ClassVar,
+    ForwardRef,
+    Generic,
+    Literal,
+    ParamSpec,
+    Self,
+    TypeAliasType,
+    TypeVar,
+    TypeVarTuple,
+    _GenericAlias,  # pyright: ignore
+    get_args,
+    get_origin,
+    get_type_hints,
+    is_typeddict,
+    overload,
+)
+
+from haiway import types as haiway_types
+from haiway.types import MISSING, Missing
+
+__all__ = [
+    "AttributeAnnotation",
+    "attribute_annotations",
+    "resolve_attribute_annotation",
+]
+
+
+class AttributeAnnotation:
+    def __init__(
+        self,
+        *,
+        origin: Any,
+        arguments: Sequence[Any] | None = None,
+        required: bool = True,
+        extra: Mapping[str, Any] | None = None,
+    ) -> None:
+        self.origin: Any = origin
+        self.arguments: Sequence[Any]
+        if arguments is None:
+            self.arguments = ()
+
+        else:
+            self.arguments = arguments
+
+        self.required: bool = required
+        self.extra: Mapping[str, Any]
+        if extra is None:
+            self.extra = {}
+
+        else:
+            self.extra = extra
+
+    def update_required(
+        self,
+        required: bool,
+        /,
+    ) -> Self:
+        self.required = self.required and required
+
+        return self
+
+    def __str__(self) -> str:
+        if alias := self.extra.get("TYPE_ALIAS"):
+            return alias
+
+        origin_str: str = getattr(self.origin, "__name__", str(self.origin))
+        arguments_str: str
+        if self.arguments:
+            arguments_str = "[" + ", ".join(str(arg) for arg in self.arguments) + "]"
+
+        else:
+            arguments_str = ""
+
+        if module := getattr(self.origin, "__module__", None):
+            return f"{module}.{origin_str}{arguments_str}"
+
+        else:
+            return f"{origin_str}{arguments_str}"
+
+
+def attribute_annotations(
+    cls: type[Any],
+    /,
+    type_parameters: Mapping[str, Any],
+) -> Mapping[str, AttributeAnnotation]:
+    self_annotation = AttributeAnnotation(
+        origin=cls,
+        # ignore arguments here, State (and draive.DataModel) will have them resolved at this stage
+        arguments=[],
+    )
+
+    # ignore args_keys here, State (and draive.DataModel) will have them resolved at this stage
+    recursion_guard: MutableMapping[str, AttributeAnnotation] = {
+        _recursion_key(cls, default=str(self_annotation)): self_annotation
+    }
+
+    attributes: dict[str, AttributeAnnotation] = {}
+    for key, annotation in get_type_hints(cls, localns={cls.__name__: cls}).items():
+        # do not include private or special items
+        if key.startswith("_"):
+            continue
+
+        # do not include ClassVars
+        if (get_origin(annotation) or annotation) is ClassVar:
+            continue
+
+        attributes[key] = resolve_attribute_annotation(
+            annotation,
+            type_parameters=type_parameters,
+            module=cls.__module__,
+            self_annotation=self_annotation,
+            recursion_guard=recursion_guard,
+        )
+
+    return attributes
+
+
+def _resolve_none(
+    annotation: Any,
+) -> AttributeAnnotation:
+    return AttributeAnnotation(origin=NoneType)
+
+
+def _resolve_missing(
+    annotation: Any,
+) -> AttributeAnnotation:
+    # special case - attributes marked as missing are not required
+    # Missing does not work properly within TypedDict though
+    return AttributeAnnotation(
+        origin=Missing,
+        required=False,
+    )
+
+
+def _resolve_literal(
+    annotation: Any,
+) -> AttributeAnnotation:
+    return AttributeAnnotation(
+        origin=Literal,
+        arguments=get_args(annotation),
+    )
+
+
+def _resolve_forward_ref(
+    annotation: ForwardRef | str,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    forward_ref: ForwardRef
+    match annotation:
+        case str() as string:
+            forward_ref = ForwardRef(string, module=module)
+
+        case reference:
+            forward_ref = reference
+
+    if evaluated := forward_ref._evaluate(
+        globalns=None,
+        localns=None,
+        recursive_guard=frozenset(),
+    ):
+        return resolve_attribute_annotation(
+            evaluated,
+            type_parameters=type_parameters,
+            module=module,
+            self_annotation=self_annotation,
+            recursion_guard=recursion_guard,
+        )
+
+    else:
+        raise RuntimeError(f"Cannot resolve annotation of {annotation}")
+
+
+def _resolve_generic_alias(  # noqa: PLR0911, PLR0912
+    annotation: GenericAlias,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    match get_origin(annotation):
+        case TypeAliasType() as alias:  # pyright: ignore[reportUnnecessaryComparison]
+            return _resolve_type_alias(
+                alias,
+                type_parameters={
+                    # verify if we should pass all parameters
+                    param.__name__: get_args(annotation)[idx]
+                    for idx, param in enumerate(alias.__type_params__)
+                },
+                module=module,
+                self_annotation=self_annotation,
+                recursion_guard=recursion_guard,
+            )
+
+        case origin if issubclass(origin, Generic):
+            match origin.__class_getitem__(  # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue]
+                tuple(
+                    type_parameters.get(
+                        arg.__name__,
+                        arg.__bound__ or Any,
+                    )
+                    if isinstance(arg, TypeVar)
+                    else arg
+                    for arg in get_args(annotation)
+                )
+            ):
+                case GenericAlias() as generic_alias:
+                    resolved_attribute = AttributeAnnotation(origin=generic_alias.__origin__)
+                    if recursion_key := _recursion_key(generic_alias):
+                        if recursive := recursion_guard.get(recursion_key):
+                            return recursive
+
+                        else:
+                            recursion_guard[recursion_key] = resolved_attribute
+
+                    resolved_attribute.arguments = [
+                        resolve_attribute_annotation(
+                            argument,
+                            type_parameters=type_parameters,
+                            module=module,
+                            self_annotation=self_annotation,
+                            recursion_guard=recursion_guard,
+                        )
+                        for argument in get_args(generic_alias)
+                    ]
+
+                    return resolved_attribute
+
+                # use resolved type if it is not an alias again
+                case resolved:  # pyright: ignore
+                    resolved_attribute = AttributeAnnotation(origin=resolved)
+
+                    if recursion_key := _recursion_key(origin):
+                        if recursive := recursion_guard.get(recursion_key):
+                            return recursive
+
+                        else:
+                            recursion_guard[recursion_key] = resolved_attribute
+
+                    resolved_attribute.arguments = [
+                        resolve_attribute_annotation(
+                            argument,
+                            type_parameters=type_parameters,
+                            module=module,
+                            self_annotation=self_annotation,
+                            recursion_guard=recursion_guard,
+                        )
+                        for argument in get_args(annotation)
+                    ]
+
+                    return resolved_attribute
+
+        case origin:
+            resolved_attribute = AttributeAnnotation(origin=origin)
+
+            if recursion_key := _recursion_key(origin):
+                if recursive := recursion_guard.get(recursion_key):
+                    return recursive
+
+            resolved_attribute.arguments = [
+                resolve_attribute_annotation(
+                    argument,
+                    type_parameters=type_parameters,
+                    module=module,
+                    self_annotation=self_annotation,
+                    recursion_guard=recursion_guard,
+                )
+                for argument in get_args(annotation)
+            ]
+
+            return resolved_attribute
+
+
+def _resolve_special_generic_alias(
+    annotation: Any,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    origin: type[Any] = get_origin(annotation)
+    resolved_attribute = AttributeAnnotation(origin=origin)
+
+    if recursion_key := _recursion_key(origin):
+        if recursive := recursion_guard.get(recursion_key):
+            return recursive
+
+        else:
+            recursion_guard[recursion_key] = resolved_attribute
+
+    resolved_attribute.arguments = [
+        resolve_attribute_annotation(
+            argument,
+            type_parameters=type_parameters,
+            module=module,
+            self_annotation=self_annotation,
+            recursion_guard=recursion_guard,
+        )
+        for argument in get_args(annotation)
+    ]
+
+    return resolved_attribute
+
+
+def _resolve_type_alias(
+    annotation: TypeAliasType,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    resolved_attribute = AttributeAnnotation(origin=MISSING)
+
+    if recursion_key := _recursion_key(annotation):
+        if recursive := recursion_guard.get(recursion_key):
+            return recursive
+
+        else:
+            recursion_guard[recursion_key] = resolved_attribute
+
+    resolved: AttributeAnnotation = resolve_attribute_annotation(
+        annotation.__value__,
+        module=annotation.__module__ or module,
+        type_parameters=type_parameters,
+        self_annotation=self_annotation,
+        recursion_guard=recursion_guard,
+    )
+
+    resolved_attribute.origin = resolved.origin
+    resolved_attribute.arguments = resolved.arguments
+    resolved_attribute.extra = {
+        **resolved.extra,
+        "TYPE_ALIAS": annotation.__name__,
+    }
+    resolved_attribute.required = resolved.required
+
+    return resolved_attribute
+
+
+def _resolve_type_var(
+    annotation: TypeVar,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    return resolve_attribute_annotation(
+        type_parameters.get(
+            annotation.__name__,
+            # use bound as default or Any otherwise
+            annotation.__bound__ or Any,
+        ),
+        module=module,
+        type_parameters=type_parameters,
+        self_annotation=self_annotation,
+        recursion_guard=recursion_guard,
+    )
+
+
+def _resolve_type_union(
+    annotation: UnionType,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    arguments: Sequence[AttributeAnnotation] = [
+        resolve_attribute_annotation(
+            argument,
+            type_parameters=type_parameters,
+            module=module,
+            self_annotation=self_annotation,
+            recursion_guard=recursion_guard,
+        )
+        for argument in get_args(annotation)
+    ]
+    return AttributeAnnotation(
+        origin=UnionType,  # pyright: ignore[reportArgumentType]
+        arguments=arguments,
+        required=all(argument.required for argument in arguments),
+    )
+
+
+def _resolve_callable(
+    annotation: Any,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    return AttributeAnnotation(
+        origin=Callable,
+        arguments=[
+            resolve_attribute_annotation(
+                argument,
+                type_parameters=type_parameters,
+                module=module,
+                self_annotation=self_annotation,
+                recursion_guard=recursion_guard,
+            )
+            for argument in get_args(annotation)
+        ],
+    )
+
+
+def _resolve_type_box(
+    annotation: Any,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    return resolve_attribute_annotation(
+        get_args(annotation)[0],
+        type_parameters=type_parameters,
+        module=module,
+        self_annotation=self_annotation,
+        recursion_guard=recursion_guard,
+    )
+
+
+def _resolve_type_not_required(
+    annotation: Any,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    return resolve_attribute_annotation(
+        get_args(annotation)[0],
+        type_parameters=type_parameters,
+        module=module,
+        self_annotation=self_annotation,
+        recursion_guard=recursion_guard,
+    ).update_required(False)
+
+
+def _resolve_type_optional(
+    annotation: Any,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    return AttributeAnnotation(
+        origin=UnionType,  # pyright: ignore[reportArgumentType]
+        arguments=[
+            resolve_attribute_annotation(
+                get_args(annotation)[0],
+                type_parameters=type_parameters,
+                module=module,
+                self_annotation=self_annotation,
+                recursion_guard=recursion_guard,
+            ),
+            AttributeAnnotation(origin=NoneType),
+        ],
+    )
+
+
+def _resolve_type_typeddict(
+    annotation: Any,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    resolved_attribute = AttributeAnnotation(origin=annotation)
+
+    if recursion_key := _recursion_key(annotation):
+        if recursive := recursion_guard.get(recursion_key):
+            return recursive
+
+        else:
+            recursion_guard[recursion_key] = resolved_attribute
+
+    resolved_attribute.arguments = [
+        resolve_attribute_annotation(
+            argument,
+            type_parameters=type_parameters,
+            module=module,
+            self_annotation=self_annotation,
+            recursion_guard=recursion_guard,
+        )
+        for argument in get_args(annotation)
+    ]
+
+    attributes: dict[str, AttributeAnnotation] = {}
+    for key, element in get_type_hints(
+        annotation,
+        localns={annotation.__name__: annotation},
+    ).items():
+        attributes[key] = resolve_attribute_annotation(
+            element,
+            type_parameters=type_parameters,
+            module=getattr(annotation, "__module__", module),
+            self_annotation=resolved_attribute,
+            recursion_guard=recursion_guard,
+        ).update_required(key in annotation.__required_keys__)
+    resolved_attribute.extra = {
+        "attributes": attributes,
+        "required": annotation.__required_keys__,
+    }
+    return resolved_attribute
+
+
+def _resolve_type(
+    annotation: Any,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    if recursion_key := _recursion_key(annotation):
+        if recursive := recursion_guard.get(recursion_key):
+            return recursive
+
+        # not updating recursion guard here - it might be a builtin type
+
+    return AttributeAnnotation(
+        origin=annotation,
+        arguments=[
+            resolve_attribute_annotation(
+                argument,
+                type_parameters=type_parameters,
+                module=module,
+                self_annotation=self_annotation,
+                recursion_guard=recursion_guard,
+            )
+            for argument in get_args(annotation)
+        ],
+    )
+
+
+def resolve_attribute_annotation(  # noqa: C901, PLR0911, PLR0912
+    annotation: Any,
+    /,
+    module: str,
+    type_parameters: Mapping[str, Any],
+    self_annotation: AttributeAnnotation | None,
+    recursion_guard: MutableMapping[str, AttributeAnnotation],
+) -> AttributeAnnotation:
+    match get_origin(annotation) or annotation:
+        case types.NoneType | None:
+            return _resolve_none(
+                annotation=annotation,
+            )
+
+        case haiway_types.Missing:
+            return _resolve_missing(
+                annotation=annotation,
+            )
+
+        case types.UnionType | typing.Union:
+            return _resolve_type_union(
+                annotation,
+                module=module,
+                type_parameters=type_parameters,
+                self_annotation=self_annotation,
+                recursion_guard=recursion_guard,
+            )
+
+        case typing.Literal:
+            return _resolve_literal(annotation)
+
+        case typeddict if is_typeddict(typeddict):
+            return _resolve_type_typeddict(
+                typeddict,
+                module=module,
+                type_parameters=type_parameters,
+                self_annotation=self_annotation,
+                recursion_guard=recursion_guard,
+            )
+
+        case typing.Callable:  # pyright: ignore
+            return _resolve_callable(
+                annotation,
+                module=module,
+                type_parameters=type_parameters,
+                self_annotation=self_annotation,
+                recursion_guard=recursion_guard,
+            )
+
+        case typing.Annotated | typing.Final | typing.Required:
+            return _resolve_type_box(
+                annotation,
+                module=module,
+                type_parameters=type_parameters,
+                self_annotation=self_annotation,
+                recursion_guard=recursion_guard,
+            )
+
+        case typing.NotRequired:
+            return _resolve_type_not_required(
+                annotation,
+                module=module,
+                type_parameters=type_parameters,
+                self_annotation=self_annotation,
+                recursion_guard=recursion_guard,
+            )
+
+        case typing.Optional:  # optional is a Union[Value, None]
+            return _resolve_type_optional(
+                annotation,
+                module=module,
+                type_parameters=type_parameters,
+                self_annotation=self_annotation,
+                recursion_guard=recursion_guard,
+            )
+
+        case typing.Self:  # pyright: ignore
+            if self_annotation:
+                return self_annotation
+
+            else:
+                raise RuntimeError(f"Unresolved Self annotation: {annotation}")
+
+        case _:
+            match annotation:
+                case str() | ForwardRef():
+                    return _resolve_forward_ref(
+                        annotation,
+                        module=module,
+                        type_parameters=type_parameters,
+                        self_annotation=self_annotation,
+                        recursion_guard=recursion_guard,
+                    )
+
+                case GenericAlias():
+                    return _resolve_generic_alias(
+                        annotation,
+                        module=module,
+                        type_parameters=type_parameters,
+                        self_annotation=self_annotation,
+                        recursion_guard=recursion_guard,
+                    )
+
+                case _GenericAlias():
+                    return _resolve_special_generic_alias(
+                        annotation,
+                        module=module,
+                        type_parameters=type_parameters,
+                        self_annotation=self_annotation,
+                        recursion_guard=recursion_guard,
+                    )
+
+                case TypeAliasType():
+                    return _resolve_type_alias(
+                        annotation,
+                        module=module,
+                        type_parameters=type_parameters,
+                        self_annotation=self_annotation,
+                        recursion_guard=recursion_guard,
+                    )
+
+                case TypeVar():
+                    return _resolve_type_var(
+                        annotation,
+                        module=module,
+                        type_parameters=type_parameters,
+                        self_annotation=self_annotation,
+                        recursion_guard=recursion_guard,
+                    )
+
+                case ParamSpec():
+                    raise NotImplementedError(f"Unresolved ParamSpec annotation: {annotation}")
+
+                case TypeVarTuple():
+                    raise NotImplementedError(f"Unresolved TypeVarTuple annotation: {annotation}")
+
+                case _:  # finally use whatever there was
+                    return _resolve_type(
+                        annotation,
+                        module=module,
+                        type_parameters=type_parameters,
+                        self_annotation=self_annotation,
+                        recursion_guard=recursion_guard,
+                    )
+
+
+@overload
+def _recursion_key(
+    annotation: Any,
+    /,
+) -> str | None: ...
+
+
+@overload
+def _recursion_key(
+    annotation: Any,
+    /,
+    default: str,
+) -> str: ...
+
+
+def _recursion_key(
+    annotation: Any,
+    /,
+    default: str | None = None,
+) -> str | None:
+    args_suffix: str
+    if arguments := get_args(annotation):
+        arguments_string: str = ", ".join(
+            _recursion_key(
+                argument,
+                default="?",
+            )
+            for argument in arguments
+        )
+        args_suffix = f"[{arguments_string}]"
+
+    else:
+        args_suffix = ""
+
+    if qualname := getattr(annotation, "__qualname__", None):
+        return qualname + args_suffix
+
+    module_prefix: str
+    if module := getattr(annotation, "__module__", None):
+        module_prefix = module + "."
+
+    else:
+        module_prefix = ""
+
+    if name := getattr(annotation, "__name__", None):
+        return module_prefix + name + args_suffix
+
+    return default