Try fixing Python 3.10 compat

Author: saulshanabrook

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/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)
@@ -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:
@@ -318,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/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]

python/tests/test_high_level.py

@@ -5,7 +5,7 @@
 import pathlib
 from copy import copy
 from fractions import Fraction
-from typing import ClassVar, TypeAlias
+from typing import ClassVar, TypeAlias, TypeVar
 
 import pytest
 
@@ -17,6 +17,7 @@
     MethodRef,
     TypedExprDecl,
 )
+from egglog.version_compat import BEFORE_3_11
 
 
 class TestExprStr:
@@ -793,10 +794,11 @@ class E(Expr):
         @function(cost=10)
         def __init__(self) -> None: ...
 
-    with pytest.raises(
-        ValueError,
-        match="Inside of classes, wrap methods with the `method` decorator, not `function`\nError processing E.__init__",
-    ):
+    match = "Inside of classes, wrap methods with the `method` decorator, not `function`"
+    # If we are after 3 11 we have context included
+    if not BEFORE_3_11:
+        match += "\nError processing E.__init__"
+    with pytest.raises(ValueError, match=match):
         E()
 
 
@@ -858,6 +860,19 @@ def __eq__(self, other: B) -> B: ...  # type: ignore[override]
     assert not isinstance(B() == B(), Fact)
 
 
+T = TypeVar("T")
+
+
+def test_type_param_sub():
+    """
+    Verify that type substituion works properly, by comparing string version.
+
+    Comparing actual versions is always false if they are no the same object for unions
+    """
+    V = Vec[T] | int
+    assert str(V[Unit]) == str(Vec[Unit] | int)  # type: ignore[misc]
+
+
 EXAMPLE_FILES = list((pathlib.Path(__file__).parent / "../egglog/examples").glob("*.py"))