typingkit 0.2.2__py3-none-any.whl

typingkit/__init__.py

@@ -0,0 +1,11 @@
+"""
+Typed NumPy
+=======
+"""
+# src/typingkit/__init__.py
+
+import numpy
+
+__all__ = [
+    "numpy",
+]

typingkit/_typed/__init__.py

@@ -0,0 +1,11 @@
+"""
+Typed NumPy core
+=======
+"""
+# src/typingkit/_typed/__init__.py
+
+from typingkit._typed.ndarray import TypedNDArray
+
+__all__ = [
+    "TypedNDArray",
+]

typingkit/_typed/_debug.py

@@ -0,0 +1,37 @@
+"""
+Debugging utilities
+=======
+"""
+# src/typingkit/_typed/debug.py
+
+from typing import Any, TypeVar, get_args, get_origin
+
+from rich.tree import Tree
+
+
+def diagnostic(obj: Any, pfx: str | None = None) -> Tree:
+    _pfx = f"[bold cyan]{pfx}[/] " if pfx else ""
+    tree = Tree(
+        f"{_pfx}[yellow]obj[/]=[green]{obj!r}[/], "
+        f"[yellow]type[/]=[magenta]{type(obj).__name__}[/]"
+    )
+    match obj:
+        case tuple():
+            for x in obj:  # type: ignore
+                tree.add(diagnostic(x))
+
+        case TypeVar():
+            tree.add(diagnostic(obj.__bound__, "__bound__:"))
+            tree.add(diagnostic(obj.__constraints__, "__constraints__:"))
+            tree.add(diagnostic(obj.__default__, "__default__:"))
+            tree.add(diagnostic(obj.__covariant__, "__covariant__:"))
+            tree.add(diagnostic(obj.__contravariant__, "__contravariant__:"))
+
+        # GenericAlias | Literal | UnionType
+        case _ if (origin := get_origin(obj)) is not None:
+            tree.add(diagnostic(origin, "origin:"))
+            tree.add(diagnostic(get_args(obj), "args:"))
+
+        case _:
+            pass
+    return tree

typingkit/_typed/context.py

@@ -0,0 +1,203 @@
+"""
+Context binding
+=======
+Manages TypeVar binding contexts for shape validation.
+"""
+# src/typingkit/_typed/context.py
+
+# pyright: reportPrivateUsage = false
+
+import inspect
+from contextvars import ContextVar
+from functools import wraps
+from typing import (
+    Any,
+    Callable,
+    Concatenate,
+    ParamSpec,
+    TypeVar,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from typingkit._typed.ndarray import (
+    DimensionError,
+    _validate_shape,
+    _validate_shape_against_contexts,
+)
+
+# Context variables
+
+_class_typevar_context = ContextVar[dict[int, dict[TypeVar, int]]](
+    "_class_typevar_context", default=dict[int, dict[TypeVar, int]]()
+)
+_method_typevar_context = ContextVar[dict[TypeVar, int]](
+    "_method_typevar_context", default=dict[TypeVar, int]()
+)
+_active_class_context = ContextVar[dict[TypeVar, int]](
+    "_active_class_context", default=dict[TypeVar, int]()
+)
+
+T = TypeVar("T")
+P = ParamSpec("P")  # ParamSpec for function parameters
+R = TypeVar("R")  # TypeVar for return type
+
+
+def _extract_shape_dims(annotation: Any) -> tuple[Any, ...] | None:
+    """Return shape dimension spec tuple[...] from TypedNDArray annotation."""
+
+    origin = get_origin(annotation)
+    if origin is None:
+        return None
+
+    args = get_args(annotation)
+    if not args:
+        return None
+
+    shape_spec = args[0]
+    if get_origin(shape_spec) is not tuple:
+        return None
+
+    return get_args(shape_spec)
+
+
+def _validate_and_bind(
+    *,
+    shape_dims: tuple[Any, ...],
+    actual_shape: tuple[int, ...],
+    owner_cls: type,
+    func_name: str,
+    param_name: str | None,
+    class_context: dict[TypeVar, int],
+    method_context: dict[TypeVar, int],
+) -> None:
+    """
+    Performs:
+    1. Structural validation (_validate_shape)
+    2. TypeVar binding + consistency checks
+    """
+
+    # Structural validation (Literal, int, rank, repeated TypeVar)
+    _validate_shape(shape_dims, actual_shape)
+
+    for dim_idx, dim in enumerate(shape_dims):
+        if not isinstance(dim, TypeVar):
+            continue
+        if dim_idx >= len(actual_shape):
+            continue
+
+        actual_dim = actual_shape[dim_idx]
+        is_class_level = _is_class_level_typevar(dim, owner_cls)
+        context = class_context if is_class_level else method_context
+
+        if dim in context:
+            expected_dim = context[dim]
+            if actual_dim != expected_dim:
+                level = "class" if is_class_level else "method"
+                location = (
+                    f"parameter `{param_name}`"
+                    if param_name is not None
+                    else "return value"
+                )
+                raise DimensionError(
+                    f"In {func_name}(...), {location} "
+                    f"dimension {dim_idx} [{dim}] "
+                    f"expected {expected_dim} ({level}-level binding), "
+                    f"got {actual_dim}"
+                )
+        else:
+            context[dim] = actual_dim
+
+
+def _is_class_level_typevar(typevar: TypeVar, owner_cls: type) -> bool:
+    """Check if a TypeVar is bound at class level vs method level."""
+    cls_params = getattr(owner_cls, "__parameters__", ())
+    return typevar in cls_params
+
+
+def _get_instance_class_context(instance: Any) -> dict[TypeVar, int]:
+    """Get or create the class-level TypeVar binding context for an instance."""
+    ctx = _class_typevar_context.get()
+    instance_id = id(instance)
+    if instance_id not in ctx:
+        ctx = ctx.copy()
+        ctx[instance_id] = {}
+        _class_typevar_context.set(ctx)
+    return ctx[instance_id]
+
+
+def enforce_shapes(
+    func: Callable[Concatenate[T, P], R],
+) -> Callable[Concatenate[T, P], R]:
+    """
+    Decorator to automatically validate TypeVar shape bindings.
+    - Class-level TypeVars (from Generic[T]) are bound per-instance, persist across calls
+    - Method-level TypeVars are validated per-call only, local to single invocation
+    - Validates both parameter and return types
+    """
+
+    @wraps(func)
+    def wrapper(self: T, *args: P.args, **kwargs: P.kwargs) -> R:
+        hints = get_type_hints(func, include_extras=True)
+        sig = inspect.signature(func)
+        bound_args = sig.bind(self, *args, **kwargs)
+        bound_args.apply_defaults()
+
+        owner_cls = self.__class__
+        class_context = _get_instance_class_context(self)
+        method_context = dict[TypeVar, int]()
+
+        # Validate arguments
+        for param_name, param_value in bound_args.arguments.items():
+            if param_name == "self":
+                continue
+            if param_name not in hints:
+                continue
+            if not hasattr(param_value, "shape"):
+                continue
+
+            shape_dims = _extract_shape_dims(hints[param_name])
+            if shape_dims is None:
+                continue
+
+            _validate_and_bind(
+                shape_dims=shape_dims,
+                actual_shape=param_value.shape,
+                owner_cls=owner_cls,
+                func_name=func.__name__,
+                param_name=param_name,
+                class_context=class_context,
+                method_context=method_context,
+            )
+
+        # Execute function with active contexts
+        method_token = _method_typevar_context.set(method_context)
+        class_token = _active_class_context.set(class_context)
+        try:
+            result = func(self, *args, **kwargs)
+        finally:
+            _method_typevar_context.reset(method_token)
+            _active_class_context.reset(class_token)
+
+        # Validate return
+        if "return" in hints and result is not None and hasattr(result, "shape"):
+            shape_dims = _extract_shape_dims(hints["return"])
+            actual_shape = getattr(result, "shape")
+            if shape_dims is not None:
+                _validate_and_bind(
+                    shape_dims=shape_dims,
+                    actual_shape=actual_shape,
+                    owner_cls=owner_cls,
+                    func_name=func.__name__,
+                    param_name=None,
+                    class_context=class_context,
+                    method_context=method_context,
+                )
+
+                # context-aware validation (cross-call)
+                _validate_shape_against_contexts(shape_dims, actual_shape)
+
+        return result
+
+    return wrapper

typingkit/_typed/dimexpr.py

@@ -0,0 +1,154 @@
+"""
+DimExpr
+=======
+"""
+# src/typingkit/_typed/dimexpr.py
+
+import math
+from typing import (
+    Any,
+    Generic,
+    Literal,
+    NoReturn,
+    TypeAlias,
+    TypeAliasType,
+    TypeVar,
+    get_args,
+    get_origin,
+)
+
+Arg1 = TypeVar("Arg1", bound=int)
+Arg2 = TypeVar("Arg2", bound=int)
+
+
+## Base
+
+
+class DimExpr(int):
+    def __new__(cls, *args: Any) -> NoReturn:
+        raise TypeError("Shape expressions cannot be instantiated")
+
+    @classmethod
+    def expr(cls, args: tuple[Any, ...], /) -> int:
+        raise NotImplementedError
+
+
+class UnaryOp(Generic[Arg1], DimExpr):
+    @classmethod
+    def expr(cls, args: tuple[Arg1], /) -> int:
+        raise NotImplementedError
+
+
+class BinaryOp(Generic[Arg1, Arg2], DimExpr):
+    @classmethod
+    def expr(cls, args: tuple[Arg1, Arg2], /) -> int:
+        raise NotImplementedError
+
+
+## Core operations
+
+
+class Neg(UnaryOp[Arg1]):
+    @classmethod
+    def expr(cls, args: tuple[Arg1], /) -> int:
+        return -args[0]
+
+
+class Add(BinaryOp[Arg1, Arg2]):
+    @classmethod
+    def expr(cls, args: tuple[Arg1, Arg2], /) -> int:
+        return args[0] + args[1]
+
+
+class Mul(BinaryOp[Arg1, Arg2]):
+    @classmethod
+    def expr(cls, args: tuple[Arg1, Arg2], /) -> int:
+        return args[0] * args[1]
+
+
+class Pow(BinaryOp[Arg1, Arg2]):
+    @classmethod
+    def expr(cls, args: tuple[Arg1, Arg2], /) -> int:
+        return args[0] ** args[1]
+
+
+## Additional operations
+
+# Through TypeAliases
+Sub: TypeAlias = Add[Arg1, Neg[Arg2]]
+PlusOne: TypeAlias = Add[Arg1, Literal[1]]
+
+# Through PEP-695 style TypeAliasTypes
+type MinusOne[Arg1: int] = Sub[Arg1, Literal[1]]
+
+
+# Through Subclassing
+class Squared(Mul[Arg1, Arg1]): ...
+
+
+# Custom operations
+# Can also directly define with a `expr`, with some custom logic, say
+class Cubed(UnaryOp[Arg1]):
+    @classmethod
+    def expr(cls, args: tuple[Arg1], /) -> int:
+        return args[0] ** 3
+
+
+class Log(BinaryOp[Arg1, Arg2]):
+    @classmethod
+    def expr(cls, args: tuple[Arg1, Arg2], /) -> int:
+        x = math.log(args[0], args[1])
+        xi = round(x)
+        if math.isclose(x, xi, rel_tol=0, abs_tol=1e-12):
+            return xi
+        raise TypeError("Invalid dimension. Not an integer.")
+
+
+Log2 = Log[Arg1, Literal[2]]
+Log10 = Log[Arg1, Literal[10]]
+
+
+## Evaluation
+
+
+def _resolve_dim(tp: Any) -> Any:
+    # type[Any] / type[int] / EllipsisType / TypeVar
+    if tp is Any or tp is int or tp is Ellipsis or isinstance(tp, TypeVar):
+        return tp
+
+    origin = get_origin(tp)
+
+    # Literal[N]
+    if origin is Literal:
+        # [TODO] How should we handle multiple args case?
+        return get_args(tp)[0]
+
+    # TypeAliasType
+    if isinstance(origin, TypeAliasType):
+        return _resolve_dim(origin.__value__[get_args(tp)])
+
+    # DimExpr
+    if origin and issubclass(origin, DimExpr):
+        args = get_args(tp)
+
+        # Case 1: class defines its own expr
+        if "expr" in origin.__dict__:
+            values = tuple(_resolve_dim(arg) for arg in args)
+            if not all(isinstance(v, int) for v in values):
+                return origin[values]  # pyright: ignore[reportInvalidTypeArguments]
+            return origin.expr(values)
+
+        # Case 2: subclassing
+        if hasattr(origin, "__orig_bases__"):
+            for base in getattr(origin, "__orig_bases__"):
+                base_origin = get_origin(base)
+                if base_origin and issubclass(base_origin, DimExpr):
+                    base_args = get_args(base)
+
+                    parameters = getattr(origin, "__parameters__", ())
+                    param_map = dict(zip(parameters, args))
+                    substituted = tuple(param_map.get(a, a) for a in base_args)
+
+                    return _resolve_dim(base_origin[substituted])  # pyright: ignore[reportInvalidTypeArguments]
+
+    raise TypeError(f"Cannot evaluate {tp}")

typingkit/_typed/factory.py

@@ -0,0 +1,71 @@
+"""
+Shape factory for TypedNDArray
+=======
+"""
+# src/typingkit/_typed/factory.py
+
+# pyright: reportPrivateUsage = false
+
+from types import GenericAlias
+from typing import Any, Generic, TypeVar
+
+import numpy as np
+import numpy.typing as npt
+
+from typingkit._typed.ndarray import TypedNDArray, _AnyShape
+
+_ShapeT = TypeVar("_ShapeT", bound=_AnyShape)
+_ScalarT = TypeVar("_ScalarT", bound=np.generic)
+_NewScalarT = TypeVar("_NewScalarT", bound=np.generic)
+
+
+class _ShapeBuilder(Generic[_ShapeT, _ScalarT]):
+    __slots__ = ("_shape_spec", "_dtype_spec")
+
+    def __init__(self, shape_spec: GenericAlias, dtype_spec: GenericAlias):
+        self._shape_spec = shape_spec
+        self._dtype_spec = dtype_spec
+
+    def dtype(self, dtype: type[_NewScalarT]) -> "_ShapeBuilder[_ShapeT, _NewScalarT]":
+        return _ShapeBuilder(self._shape_spec, GenericAlias(np.dtype, (dtype,)))
+
+    def __call__(
+        self, object: npt.ArrayLike
+    ) -> TypedNDArray[_ShapeT, np.dtype[_ScalarT]]:
+        dtype = self._dtype_spec.__args__[0]
+        dtype = None if dtype is Any else dtype
+        return TypedNDArray[self._shape_spec](object, dtype)  # type: ignore
+
+    def __repr__(self) -> str:
+        return f"TypedNDArrayFactory[{self._shape_spec.__args__}, {self._dtype_spec.__args__[0]}]"
+
+    ## Convenience helpers
+
+    @property
+    def float64(self) -> "_ShapeBuilder[_ShapeT, np.float64]":
+        return self.dtype(np.float64)
+
+    @property
+    def float32(self) -> "_ShapeBuilder[_ShapeT, np.float32]":
+        return self.dtype(np.float32)
+
+    @property
+    def int64(self) -> "_ShapeBuilder[_ShapeT, np.int64]":
+        return self.dtype(np.int64)
+
+
+class _ShapeFactory:
+    def __getitem__(self, dims: _ShapeT) -> _ShapeBuilder[_ShapeT, Any]:
+        if not isinstance(dims, tuple):  # pyright: ignore[reportUnnecessaryIsInstance]
+            dims = (dims,)
+
+        shape_spec = GenericAlias(tuple, dims)
+        dtype_spec = GenericAlias(np.dtype, (Any,))
+        return _ShapeBuilder(shape_spec, dtype_spec)
+
+    # Prefer [...] rather than (...), but this is also provided
+    def __call__(self, dims: _ShapeT) -> _ShapeBuilder[_ShapeT, Any]:
+        return self.__getitem__(dims)
+
+
+Shaped = _ShapeFactory()

typingkit/_typed/generics.py

@@ -0,0 +1,50 @@
+"""
+Generics
+=======
+"""
+# src/typingkit/_typed/generics.py
+
+from types import GenericAlias
+from typing import Any, Generic, Self, TypeVarTuple, Unpack, cast, get_args, get_origin
+
+Ts = TypeVarTuple("Ts", default=Unpack[tuple[Any, ...]])
+
+
+class RuntimeGeneric(Generic[Unpack[Ts]]):
+    @classmethod
+    def __class_getitem__(cls, item: Any, /) -> GenericAlias:
+        # [HACK] Misuses __class_getitem__
+        # See https://docs.python.org/3/reference/datamodel.html#the-purpose-of-class-getitem
+
+        try:
+            ga = cast(GenericAlias, super().__class_getitem__(item))  # type: ignore[misc]
+        except:  # noqa: E722
+            # Fallback if superclass does not implement `__class_getitem__`
+            ga = GenericAlias(cls, item)
+        return _RuntimeGenericAlias.from_generic_alias(ga)
+
+    @classmethod
+    def __pre_new__(cls, alias: GenericAlias, *args: Any, **kwargs: Any) -> Self:
+        return cls(*args, **kwargs)
+
+
+class _RuntimeGenericAlias(GenericAlias):
+    """
+    Deferred RuntimeGeneric constructor.
+    Enables progressive type specialisation, behaving like a type-level curry.
+    """
+
+    @classmethod
+    def from_generic_alias(cls, alias: GenericAlias) -> Self:
+        origin = get_origin(alias)
+        typeargs = get_args(alias)
+        return cls(origin, typeargs)
+
+    def __getitem__(self, typeargs: Any) -> Self:
+        ga = super().__getitem__(typeargs)
+        return type(self).from_generic_alias(ga)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        origin: type[RuntimeGeneric] = get_origin(self)
+        obj = origin.__pre_new__(self, *args, **kwargs)
+        return obj