egglog 11.2.0__cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

egglog/thunk.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+from typing_extensions import TypeVarTuple, Unpack
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+__all__ = ["Thunk", "split_thunk"]
+
+T = TypeVar("T")
+TS = TypeVarTuple("TS")
+V = TypeVar("V")
+
+
+def split_thunk(fn: Callable[[], tuple[T, V]]) -> tuple[Callable[[], T], Callable[[], V]]:
+    s = _Split(fn)
+    return s.left, s.right
+
+
+@dataclass
+class _Split(Generic[T, V]):
+    fn: Callable[[], tuple[T, V]]
+
+    def left(self) -> T:
+        return self.fn()[0]
+
+    def right(self) -> V:
+        return self.fn()[1]
+
+
+@dataclass
+class Thunk(Generic[T, Unpack[TS]]):
+    """
+    Cached delayed function call.
+    """
+
+    state: Resolved[T] | Unresolved[T, Unpack[TS]] | Resolving | Error
+
+    @classmethod
+    def fn(cls, fn: Callable[[Unpack[TS]], T], *args: Unpack[TS], context: str | None = None) -> Thunk[T, Unpack[TS]]:
+        """
+        Create a thunk based on some functions and some partial args.
+
+        If the function is called while it is being resolved recursively it will raise an exception.
+        """
+        return cls(Unresolved(fn, args, context))
+
+    @classmethod
+    def value(cls, value: T) -> Thunk[T]:
+        return Thunk(Resolved(value))
+
+    def __call__(self) -> T:
+        match self.state:
+            case Resolved(value):
+                return value
+            case Unresolved(fn, args, context):
+                self.state = Resolving()
+                try:
+                    res = fn(*args)
+                except Exception as e:
+                    self.state = Error(e, context)
+                    raise e from None
+                else:
+                    self.state = Resolved(res)
+                    return res
+            case Resolving():
+                msg = "Recursively resolving thunk"
+                raise ValueError(msg)
+            case Error(e):
+                raise e
+
+
+@dataclass
+class Resolved(Generic[T]):
+    value: T
+
+
+@dataclass
+class Unresolved(Generic[T, Unpack[TS]]):
+    fn: Callable[[Unpack[TS]], T]
+    args: tuple[Unpack[TS]]
+    context: str | None
+
+
+@dataclass
+class Resolving:
+    pass
+
+
+@dataclass
+class Error:
+    e: Exception
+    context: str | None

egglog/type_constraint_solver.py

@@ -0,0 +1,113 @@
+"""Provides a class for solving type constraints."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass, field
+from itertools import chain, repeat
+from typing import TYPE_CHECKING
+
+from typing_extensions import assert_never
+
+from .declarations import *
+
+if TYPE_CHECKING:
+    from collections.abc import Collection, Iterable
+
+
+__all__ = ["TypeConstraintError", "TypeConstraintSolver"]
+
+
+class TypeConstraintError(RuntimeError):
+    """Typing error when trying to infer the return type."""
+
+
+@dataclass
+class TypeConstraintSolver:
+    """
+    Given some typevars and types, solves the constraints to resolve the typevars.
+    """
+
+    _decls: Declarations = field(repr=False)
+    # Mapping of class name to mapping of bound class typevar to type
+    _cls_typevar_index_to_type: defaultdict[str, dict[ClassTypeVarRef, JustTypeRef]] = field(
+        default_factory=lambda: defaultdict(dict)
+    )
+
+    def bind_class(self, ref: JustTypeRef) -> None:
+        """
+        Bind the typevars of a class to the given types.
+        Used for a situation like Map[int, str].create().
+        """
+        name = ref.name
+        cls_typevars = self._decls.get_class_decl(name).type_vars
+        if len(cls_typevars) != len(ref.args):
+            raise TypeConstraintError(f"Mismatch of typevars {cls_typevars} and {ref}")
+        bound_typevars = self._cls_typevar_index_to_type[name]
+        for i, arg in enumerate(ref.args):
+            bound_typevars[cls_typevars[i]] = arg
+
+    def infer_arg_types(
+        self,
+        fn_args: Collection[TypeOrVarRef],
+        fn_return: TypeOrVarRef,
+        fn_var_args: TypeOrVarRef | None,
+        return_: JustTypeRef,
+        cls_name: str | None,
+    ) -> tuple[Iterable[JustTypeRef], tuple[JustTypeRef, ...] | None]:
+        """
+        Given a return type, infer the argument types. If there is a variable arg, it returns an infinite iterable.
+
+        Also returns the bound type params if the class name is passed in.
+        """
+        self.infer_typevars(fn_return, return_, cls_name)
+        arg_types: Iterable[JustTypeRef] = [self.substitute_typevars(a, cls_name) for a in fn_args]
+        if fn_var_args:
+            # Need to be generator so it can be infinite for variable args
+            arg_types = chain(arg_types, repeat(self.substitute_typevars(fn_var_args, cls_name)))
+        bound_typevars = (
+            tuple(
+                v
+                # Sort by the index of the typevar in the class
+                for _, v in sorted(
+                    self._cls_typevar_index_to_type[cls_name].items(),
+                    key=lambda kv: self._decls.get_class_decl(cls_name).type_vars.index(kv[0]),
+                )
+            )
+            if cls_name
+            else None
+        )
+        return arg_types, bound_typevars
+
+    def infer_typevars(self, fn_arg: TypeOrVarRef, arg: JustTypeRef, cls_name: str | None = None) -> None:
+        match fn_arg:
+            case TypeRefWithVars(cls_name, fn_args):
+                if cls_name != arg.name:
+                    raise TypeConstraintError(f"Expected {cls_name}, got {arg.name}")
+                for inner_fn_arg, inner_arg in zip(fn_args, arg.args, strict=True):
+                    self.infer_typevars(inner_fn_arg, inner_arg, cls_name)
+            case ClassTypeVarRef():
+                if cls_name is None:
+                    msg = "Cannot infer typevar without class name"
+                    raise RuntimeError(msg)
+
+                class_typevars = self._cls_typevar_index_to_type[cls_name]
+                if fn_arg in class_typevars:
+                    if class_typevars[fn_arg] != arg:
+                        raise TypeConstraintError(f"Expected {class_typevars[fn_arg]}, got {arg}")
+                else:
+                    class_typevars[fn_arg] = arg
+            case _:
+                assert_never(fn_arg)
+
+    def substitute_typevars(self, tp: TypeOrVarRef, cls_name: str | None = None) -> JustTypeRef:
+        match tp:
+            case ClassTypeVarRef():
+                assert cls_name is not None
+                try:
+                    return self._cls_typevar_index_to_type[cls_name][tp]
+                except KeyError as e:
+                    raise TypeConstraintError(f"Not enough bound typevars for {tp!r} in class {cls_name}") from e
+            case TypeRefWithVars(name, args):
+                return JustTypeRef(name, tuple(self.substitute_typevars(arg, cls_name) for arg in args))
+        assert_never(tp)

egglog/version_compat.py

@@ -0,0 +1,87 @@
+import collections
+import sys
+import types
+import typing
+
+BEFORE_3_11 = sys.version_info < (3, 11)
+
+__all__ = ["add_note"]
+
+
+def add_note(message: str, exc: BaseException) -> BaseException:
+    """
+    Backwards compatible add_note for Python <= 3.10
+    """
+    if BEFORE_3_11:
+        return exc
+    exc.add_note(message)
+    return exc
+
+
+# For Python version 3.10 need to monkeypatch this function so that RuntimeClass type parameters
+# will be collected as typevars
+if BEFORE_3_11:
+
+    @typing.no_type_check
+    def _collect_type_vars_monkeypatch(types_, typevar_types=None):
+        """
+        Collect all type variable contained
+        in types in order of first appearance (lexicographic order). For example::
+
+            _collect_type_vars((T, List[S, T])) == (T, S)
+        """
+        from .runtime import RuntimeClass  # noqa: PLC0415
+
+        if typevar_types is None:
+            typevar_types = typing.TypeVar
+        tvars = []
+        for t in types_:
+            if isinstance(t, typevar_types) and t not in tvars:
+                tvars.append(t)
+            # **MONKEYPATCH CHANGE HERE TO ADD RuntimeClass**
+            if isinstance(t, (typing._GenericAlias, typing.GenericAlias, types.UnionType, RuntimeClass)):
+                tvars.extend([t for t in t.__parameters__ if t not in tvars])
+        return tuple(tvars)
+
+    typing._collect_type_vars = _collect_type_vars_monkeypatch  # type: ignore[attr-defined]
+
+    @typing.no_type_check
+    @typing._tp_cache
+    def __getitem__monkeypatch(self, params):  # noqa: C901, PLR0912
+        from .runtime import RuntimeClass  # noqa: PLC0415
+
+        if self.__origin__ in (typing.Generic, typing.Protocol):
+            # Can't subscript Generic[...] or Protocol[...].
+            raise TypeError(f"Cannot subscript already-subscripted {self}")
+        if not isinstance(params, tuple):
+            params = (params,)
+        params = tuple(typing._type_convert(p) for p in params)
+        if self._paramspec_tvars and any(isinstance(t, typing.ParamSpec) for t in self.__parameters__):
+            params = typing._prepare_paramspec_params(self, params)
+        else:
+            typing._check_generic(self, params, len(self.__parameters__))
+
+        subst = dict(zip(self.__parameters__, params, strict=False))
+        new_args = []
+        for arg in self.__args__:
+            if isinstance(arg, self._typevar_types):
+                if isinstance(arg, typing.ParamSpec):
+                    arg = subst[arg]  # noqa: PLW2901
+                    if not typing._is_param_expr(arg):
+                        raise TypeError(f"Expected a list of types, an ellipsis, ParamSpec, or Concatenate. Got {arg}")
+                else:
+                    arg = subst[arg]  # noqa: PLW2901
+            # **MONKEYPATCH CHANGE HERE TO ADD RuntimeClass**
+            elif isinstance(arg, (typing._GenericAlias, typing.GenericAlias, types.UnionType, RuntimeClass)):
+                subparams = arg.__parameters__
+                if subparams:
+                    subargs = tuple(subst[x] for x in subparams)
+                    arg = arg[subargs]  # noqa: PLW2901
+            # Required to flatten out the args for CallableGenericAlias
+            if self.__origin__ == collections.abc.Callable and isinstance(arg, tuple):
+                new_args.extend(arg)
+            else:
+                new_args.append(arg)
+        return self.copy_with(tuple(new_args))
+
+    typing._GenericAlias.__getitem__ = __getitem__monkeypatch  # type: ignore[attr-defined]