Merge pull request #302 from egraphs-good/fix-python-3.10

Fix python 3.10 compat

Author: saulshanabrook

.github/workflows/CI.yml

@@ -42,8 +42,6 @@ jobs:
       - uses: astral-sh/setup-uv@v6
         with:
           enable-cache: true
-          # Run on oldest Python version to catch more errors
-          python-version: "3.10"
       - uses: dtolnay/rust-toolchain@1.79.0
       - uses: Swatinem/rust-cache@v2
       - run: uv sync --extra test --locked

docs/changelog.md

@@ -4,7 +4,8 @@ _This project uses semantic versioning_
 
 ## UNRELEASED
 
-- fix using `f64Like` when not importing star (also properly includes removal of `Callable` special case from previous release).
+- Fix using `f64Like` when not importing star (also properly includes removal of `Callable` special case from previous release).
+- Fix Python 3.10 compatibility
 
 ## 10.0.1 (2025-04-06)
 

python/egglog/declarations.py

@@ -93,7 +93,7 @@ def __egg_decls__(self) -> Declarations:
         # Catch attribute error, so that it isn't bubbled up as a missing attribute and fallbacks on `__getattr__`
         # instead raise explicitly
         except AttributeError as err:
-            msg = f"Cannot resolve declerations for {self}"
+            msg = f"Cannot resolve declarations for {self}"
             raise RuntimeError(msg) from err
 
 
@@ -308,11 +308,11 @@ class ClassTypeVarRef:
     module: str
 
     def to_just(self) -> JustTypeRef:
-        msg = "egglog does not support generic classes yet."
+        msg = f"{self}: egglog does not support generic classes yet."
         raise NotImplementedError(msg)
 
     def __str__(self) -> str:
-        return f"{self.module}.{self.name}"
+        return str(self.to_type_var())
 
     @classmethod
     def from_type_var(cls, typevar: TypeVar) -> ClassTypeVarRef:

python/egglog/egraph.py

@@ -16,7 +16,6 @@
     ClassVar,
     Generic,
     Literal,
-    Never,
     TypeAlias,
     TypedDict,
     TypeVar,
@@ -26,7 +25,7 @@
 )
 
 import graphviz
-from typing_extensions import ParamSpec, Self, Unpack, assert_never
+from typing_extensions import Never, ParamSpec, Self, Unpack, assert_never
 
 from . import bindings
 from .conversion import *
@@ -36,6 +35,7 @@
 from .pretty import pretty_decl
 from .runtime import *
 from .thunk import *
+from .version_compat import *
 
 if TYPE_CHECKING:
     from .builtins import String, Unit
@@ -169,8 +169,9 @@ def check_eq(x: BASE_EXPR, y: BASE_EXPR, schedule: Schedule | None = None, *, ad
     except bindings.EggSmolError as err:
         if display:
             egraph.display()
-        err.add_note(f"Failed:\n{eq(x).to(y)}\n\nExtracted:\n {eq(egraph.extract(x)).to(egraph.extract(y))})")
-        raise
+        raise add_note(
+            f"Failed:\n{eq(x).to(y)}\n\nExtracted:\n {eq(egraph.extract(x)).to(egraph.extract(y))})", err
+        ) from None
     return egraph
 
 
@@ -492,8 +493,7 @@ def _generate_class_decls(  # noqa: C901,PLR0912
                 reverse_args=reverse_args,
             )
         except Exception as e:
-            e.add_note(f"Error processing {cls_name}.{method_name}")
-            raise
+            raise add_note(f"Error processing {cls_name}.{method_name}", e) from None
 
         if not builtin and not isinstance(ref, InitRef) and not mutates:
             add_default_funcs.append(add_rewrite)
@@ -627,7 +627,7 @@ def _fn_decl(
     )
     decls |= merged
 
-    # defer this in generator so it doesnt resolve for builtins eagerly
+    # defer this in generator so it doesn't resolve for builtins eagerly
     args = (TypedExprDecl(tp.to_just(), VarDecl(name, False)) for name, tp in zip(arg_names, arg_types, strict=True))
     res_ref: FunctionRef | MethodRef | ClassMethodRef | PropertyRef | InitRef | UnnamedFunctionRef
     res_thunk: Callable[[], object]
@@ -671,7 +671,7 @@ def _fn_decl(
             )
         res_ref = ref
         decls.set_function_decl(ref, decl)
-        res_thunk = Thunk.fn(_create_default_value, decls, ref, fn, args, ruleset)
+        res_thunk = Thunk.fn(_create_default_value, decls, ref, fn, args, ruleset, context=f"creating {ref}")
     return res_ref, Thunk.fn(_add_default_rewrite_function, decls, res_ref, return_type, ruleset, res_thunk, subsume)
 
 
@@ -1040,8 +1040,7 @@ def _run_extract(self, expr: RuntimeExpr, n: int) -> bindings._ExtractReport:
                 bindings.ActionCommand(bindings.Extract(span(2), expr, bindings.Lit(span(2), bindings.Int(n))))
             )
         except BaseException as e:
-            e.add_note("Extracting: " + str(expr))
-            raise
+            raise add_note("Extracting: " + str(expr), e)  # noqa: B904
         extract_report = self._egraph.extract_report()
         if not extract_report:
             msg = "No extract report saved"

python/egglog/exp/array_api.py

@@ -69,6 +69,7 @@
 
 from egglog import *
 from egglog.runtime import RuntimeExpr
+from egglog.version_compat import add_note
 
 from .program_gen import *
 
@@ -1198,13 +1199,13 @@ def if_(cls, b: BooleanLike, i: NDArrayLike, j: NDArrayLike) -> NDArray: ...
 
 NDArrayLike: TypeAlias = NDArray | ValueLike | TupleValueLike
 
-converter(NDArray, IndexKey, IndexKey.ndarray)
-converter(Value, NDArray, NDArray.scalar)
+converter(NDArray, IndexKey, lambda v: IndexKey.ndarray(v))
+converter(Value, NDArray, lambda v: NDArray.scalar(v))
 # Need this if we want to use ints in slices of arrays coming from 1d arrays, but make it more expensive
 # to prefer upcasting in the other direction when we can, which is safer at runtime
 converter(NDArray, Value, lambda n: n.to_value(), 100)
-converter(TupleValue, NDArray, NDArray.vector)
-converter(TupleInt, TupleValue, TupleValue.from_tuple_int)
+converter(TupleValue, NDArray, lambda v: NDArray.vector(v))
+converter(TupleInt, TupleValue, lambda v: TupleValue.from_tuple_int(v))
 
 
 @array_api_ruleset.register
@@ -1383,8 +1384,8 @@ def int(cls, value: Int) -> IntOrTuple: ...
     def tuple(cls, value: TupleIntLike) -> IntOrTuple: ...
 
 
-converter(Int, IntOrTuple, IntOrTuple.int)
-converter(TupleInt, IntOrTuple, IntOrTuple.tuple)
+converter(Int, IntOrTuple, lambda v: IntOrTuple.int(v))
+converter(TupleInt, IntOrTuple, lambda v: IntOrTuple.tuple(v))
 
 
 class OptionalIntOrTuple(Expr, ruleset=array_api_ruleset):
@@ -1395,7 +1396,7 @@ def some(cls, value: IntOrTuple) -> OptionalIntOrTuple: ...
 
 
 converter(type(None), OptionalIntOrTuple, lambda _: OptionalIntOrTuple.none)
-converter(IntOrTuple, OptionalIntOrTuple, OptionalIntOrTuple.some)
+converter(IntOrTuple, OptionalIntOrTuple, lambda v: OptionalIntOrTuple.some(v))
 
 
 @function
@@ -1980,6 +1981,5 @@ def try_evaling(egraph: EGraph, schedule: Schedule, expr: Expr, prim_expr: Built
             extracted = egraph.extract(prim_expr)
         except BaseException as e:
             # egraph.display(n_inline_leaves=1, split_primitive_outputs=True)
-            e.add_note(f"Cannot evaluate {egraph.extract(expr)}")
-            raise
+            raise add_note(f"Cannot evaluate {egraph.extract(expr)}", e)  # noqa: B904
     return extracted.eval()  # type: ignore[attr-defined]

python/egglog/runtime.py

@@ -22,6 +22,7 @@
 from .pretty import *
 from .thunk import Thunk
 from .type_constraint_solver import *
+from .version_compat import *
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -249,8 +250,7 @@ def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable:
         try:
             cls_decl = self.__egg_decls__._classes[self.__egg_tp__.name]
         except Exception as e:
-            e.add_note(f"Error processing class {self.__egg_tp__.name}")
-            raise
+            raise add_note(f"Error processing class {self.__egg_tp__.name}", e) from None
 
         preserved_methods = cls_decl.preserved_methods
         if name in preserved_methods:
@@ -281,6 +281,9 @@ def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable:
     def __str__(self) -> str:
         return str(self.__egg_tp__)
 
+    def __repr__(self) -> str:
+        return str(self)
+
     # Make hashable so can go in Union
     def __hash__(self) -> int:
         return hash((id(self.__egg_decls_thunk__), self.__egg_tp__))
@@ -315,8 +318,7 @@ def __call__(self, *args: object, _egg_partial_function: bool = False, **kwargs:
         try:
             signature = self.__egg_decls__.get_callable_decl(self.__egg_ref__).signature
         except Exception as e:
-            e.add_note(f"Failed to find callable {self}")
-            raise
+            raise add_note(f"Failed to find callable {self}", e)  # noqa: B904
         decls = self.__egg_decls__.copy()
         # Special case function application bc we dont support variadic generics yet generally
         if signature == "fn-app":

python/egglog/thunk.py

@@ -41,13 +41,13 @@ class Thunk(Generic[T, Unpack[TS]]):
     state: Resolved[T] | Unresolved[T, Unpack[TS]] | Resolving | Error
 
     @classmethod
-    def fn(cls, fn: Callable[[Unpack[TS]], T], *args: Unpack[TS]) -> Thunk[T, Unpack[TS]]:
+    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))
+        return cls(Unresolved(fn, args, context))
 
     @classmethod
     def value(cls, value: T) -> Thunk[T]:
@@ -57,12 +57,12 @@ def __call__(self) -> T:
         match self.state:
             case Resolved(value):
                 return value
-            case Unresolved(fn, args):
+            case Unresolved(fn, args, context):
                 self.state = Resolving()
                 try:
                     res = fn(*args)
                 except Exception as e:
-                    self.state = Error(e)
+                    self.state = Error(e, context)
                     raise e from None
                 else:
                     self.state = Resolved(res)
@@ -83,6 +83,7 @@ class Resolved(Generic[T]):
 class Unresolved(Generic[T, Unpack[TS]]):
     fn: Callable[[Unpack[TS]], T]
     args: tuple[Unpack[TS]]
+    context: str | None
 
 
 @dataclass
@@ -93,3 +94,4 @@ class Resolving:
 @dataclass
 class Error:
     e: Exception
+    context: str | None

python/egglog/type_constraint_solver.py

@@ -107,7 +107,7 @@ def substitute_typevars(self, tp: TypeOrVarRef, cls_name: str | None = 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} in class {cls_name}") from 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)

python/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
+
+        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)):  # type: ignore[name-defined]
+                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
+
+        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]