egglog 10.0.1__cp312-cp312-win_amd64.whl → 11.0.0__cp312-cp312-win_amd64.whl

egglog/runtime.py

@@ -11,30 +11,32 @@ so they are not mangled by Python and can be accessed by the user.
 
 from __future__ import annotations
 
+import itertools
 import operator
+import types
 from collections.abc import Callable
-from dataclasses import dataclass, replace
+from dataclasses import InitVar, dataclass, replace
 from inspect import Parameter, Signature
 from itertools import zip_longest
-from typing import TYPE_CHECKING, TypeVar, Union, cast, get_args, get_origin
+from typing import TYPE_CHECKING, Any, TypeVar, Union, cast, get_args, get_origin
 
 from .declarations import *
 from .pretty import *
 from .thunk import Thunk
 from .type_constraint_solver import *
+from .version_compat import *
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
-    from .egraph import Fact
-
 
 __all__ = [
     "LIT_CLASS_NAMES",
-    "REFLECTED_BINARY_METHODS",
+    "NUMERIC_BINARY_METHODS",
     "RuntimeClass",
     "RuntimeExpr",
     "RuntimeFunction",
+    "define_expr_method",
     "resolve_callable",
     "resolve_type_annotation",
     "resolve_type_annotation_mutate",
@@ -45,24 +47,34 @@ UNIT_CLASS_NAME = "Unit"
 UNARY_LIT_CLASS_NAMES = {"i64", "f64", "Bool", "String"}
 LIT_CLASS_NAMES = UNARY_LIT_CLASS_NAMES | {UNIT_CLASS_NAME, "PyObject"}
 
-REFLECTED_BINARY_METHODS = {
-    "__radd__": "__add__",
-    "__rsub__": "__sub__",
-    "__rmul__": "__mul__",
-    "__rmatmul__": "__matmul__",
-    "__rtruediv__": "__truediv__",
-    "__rfloordiv__": "__floordiv__",
-    "__rmod__": "__mod__",
-    "__rpow__": "__pow__",
-    "__rlshift__": "__lshift__",
-    "__rrshift__": "__rshift__",
-    "__rand__": "__and__",
-    "__rxor__": "__xor__",
-    "__ror__": "__or__",
+# All methods which should return NotImplemented if they fail to resolve and are reflected as well
+# From https://docs.python.org/3/reference/datamodel.html#emulating-numeric-types
+
+NUMERIC_BINARY_METHODS = {
+    "__add__",
+    "__sub__",
+    "__mul__",
+    "__matmul__",
+    "__truediv__",
+    "__floordiv__",
+    "__mod__",
+    "__divmod__",
+    "__pow__",
+    "__lshift__",
+    "__rshift__",
+    "__and__",
+    "__xor__",
+    "__or__",
+    "__lt__",
+    "__le__",
+    "__gt__",
+    "__ge__",
 }
 
-# Methods that need to return real Python values not expressions
-PRESERVED_METHODS = [
+
+# Methods that need to be defined on the runtime type that holds `Expr` objects, so that they can be used as methods.
+
+TYPE_DEFINED_METHODS = {
     "__bool__",
     "__len__",
     "__complex__",
@@ -70,9 +82,15 @@ PRESERVED_METHODS = [
     "__float__",
     "__iter__",
     "__index__",
-    "__float__",
-    "__int__",
-]
+    "__call__",
+    "__getitem__",
+    "__setitem__",
+    "__delitem__",
+    "__pos__",
+    "__neg__",
+    "__invert__",
+    "__round__",
+}
 
 # Set this globally so we can get access to PyObject when we have a type annotation of just object.
 # This is the only time a type annotation doesn't need to include the egglog type b/c object is top so that would be redundant statically.
@@ -134,11 +152,48 @@ def inverse_resolve_type_annotation(decls_thunk: Callable[[], Declarations], tp:
 ##
 
 
-@dataclass
-class RuntimeClass(DelayedDeclerations):
+class BaseClassFactoryMeta(type):
+    """
+    Base metaclass for all runtime classes created by ClassFactory
+    """
+
+    def __instancecheck__(cls, instance: object) -> bool:
+        assert isinstance(cls, RuntimeClass)
+        return isinstance(instance, RuntimeExpr) and cls.__egg_tp__.name == instance.__egg_typed_expr__.tp.name
+
+
+class ClassFactory(type):
+    """
+    A metaclass for types which should create `type` objects when instantiated.
+
+    That's so that they work with `isinstance` and can be placed in `match ClassName()`.
+    """
+
+    def __call__(cls, *args, **kwargs) -> type:
+        # If we have params, don't inherit from `type` because we don't need to match against this and also
+        # this won't work with `Union[X]` because it won't look at `__parameters__` for instances of `type`.
+        if kwargs.pop("_egg_has_params", False):
+            return super().__call__(*args, **kwargs)
+        namespace: dict[str, Any] = {}
+        for m in reversed(cls.__mro__):
+            namespace.update(m.__dict__)
+        init = namespace.pop("__init__")
+        meta = types.new_class("type(RuntimeClass)", (BaseClassFactoryMeta,), {}, lambda ns: ns.update(**namespace))
+        tp = types.new_class("RuntimeClass", (), {"metaclass": meta})
+        init(tp, *args, **kwargs)
+        return tp
+
+    def __instancecheck__(cls, instance: object) -> bool:
+        return isinstance(instance, BaseClassFactoryMeta)
+
+
+@dataclass(match_args=False)
+class RuntimeClass(DelayedDeclerations, metaclass=ClassFactory):
     __egg_tp__: TypeRefWithVars
+    # True if we want `__parameters__` to be recognized by `Union`, which means we can't inherit from `type` directly.
+    _egg_has_params: InitVar[bool] = False
 
-    def __post_init__(self) -> None:
+    def __post_init__(self, _egg_has_params: bool) -> None:
         global _PY_OBJECT_CLASS, _UNSTABLE_FN_CLASS
         if (name := self.__egg_tp__.name) == "PyObject":
             _PY_OBJECT_CLASS = self
@@ -228,7 +283,7 @@ class RuntimeClass(DelayedDeclerations):
         else:
             final_args = new_args
         tp = TypeRefWithVars(self.__egg_tp__.name, final_args)
-        return RuntimeClass(Thunk.fn(Declarations.create, self, *decls_like), tp)
+        return RuntimeClass(Thunk.fn(Declarations.create, self, *decls_like), tp, _egg_has_params=True)
 
     def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable:
         if name == "__origin__" and self.__egg_tp__.args:
@@ -249,8 +304,7 @@ class RuntimeClass(DelayedDeclerations):
         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,10 +335,21 @@ class RuntimeClass(DelayedDeclerations):
     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__))
 
+    def __eq__(self, other: object) -> bool:
+        """
+        Support equality for runtime comparison of egglog classes.
+        """
+        if not isinstance(other, RuntimeClass):
+            return NotImplemented
+        return self.__egg_tp__ == other.__egg_tp__
+
     # Support unioning like types
     def __or__(self, value: type) -> object:
         return Union[self, value]  # noqa: UP007
@@ -296,6 +361,10 @@ class RuntimeClass(DelayedDeclerations):
         """
         return tuple(inverse_resolve_type_annotation(self.__egg_decls_thunk__, tp) for tp in self.__egg_tp__.args)
 
+    @property
+    def __match_args__(self) -> tuple[str, ...]:
+        return self.__egg_decls__._classes[self.__egg_tp__.name].match_args
+
 
 @dataclass
 class RuntimeFunction(DelayedDeclerations):
@@ -303,20 +372,30 @@ class RuntimeFunction(DelayedDeclerations):
     # bound methods need to store RuntimeExpr not just TypedExprDecl, so they can mutate the expr if required on self
     __egg_bound__: JustTypeRef | RuntimeExpr | None = None
 
+    def __eq__(self, other: object) -> bool:
+        """
+        Support equality for runtime comparison of egglog functions.
+        """
+        if not isinstance(other, RuntimeFunction):
+            return NotImplemented
+        return self.__egg_ref__ == other.__egg_ref__ and bool(self.__egg_bound__ == other.__egg_bound__)
+
+    def __hash__(self) -> int:
+        return hash((self.__egg_ref__, self.__egg_bound__))
+
     @property
     def __egg_ref__(self) -> CallableRef:
         return self.__egg_ref_thunk__()
 
     def __call__(self, *args: object, _egg_partial_function: bool = False, **kwargs: object) -> RuntimeExpr | None:
-        from .conversion import resolve_literal
+        from .conversion import resolve_literal  # noqa: PLC0415
 
         if isinstance(self.__egg_bound__, RuntimeExpr):
             args = (self.__egg_bound__, *args)
         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":
@@ -344,7 +423,7 @@ class RuntimeFunction(DelayedDeclerations):
         try:
             bound = py_signature.bind(*args, **kwargs)
         except TypeError as err:
-            raise TypeError(f"Failed to call {self} with args {args} and kwargs {kwargs}") from err
+            raise TypeError(f"Failed to bind arguments for {self} with args {args} and kwargs {kwargs}: {err}") from err
         del kwargs
         bound.apply_defaults()
         assert not bound.kwargs
@@ -413,9 +492,7 @@ def to_py_signature(sig: FunctionSignature, decls: Declarations, optional_args:
         Parameter(
             n,
             Parameter.POSITIONAL_OR_KEYWORD,
-            default=RuntimeExpr.__from_values__(
-                decls, TypedExprDecl(t.to_just(), d if d is not None else VarDecl(n, True))
-            )
+            default=RuntimeExpr.__from_values__(decls, TypedExprDecl(t.to_just(), d or LetRefDecl(n)))
             if d is not None or optional_args
             else Parameter.empty,
         )
@@ -426,32 +503,6 @@ def to_py_signature(sig: FunctionSignature, decls: Declarations, optional_args:
     return Signature(parameters)
 
 
-# All methods which should return NotImplemented if they fail to resolve
-# From https://docs.python.org/3/reference/datamodel.html
-PARTIAL_METHODS = {
-    "__lt__",
-    "__le__",
-    "__eq__",
-    "__ne__",
-    "__gt__",
-    "__ge__",
-    "__add__",
-    "__sub__",
-    "__mul__",
-    "__matmul__",
-    "__truediv__",
-    "__floordiv__",
-    "__mod__",
-    "__divmod__",
-    "__pow__",
-    "__lshift__",
-    "__rshift__",
-    "__and__",
-    "__xor__",
-    "__or__",
-}
-
-
 @dataclass
 class RuntimeExpr(DelayedDeclerations):
     __egg_typed_expr_thunk__: Callable[[], TypedExprDecl]
@@ -468,17 +519,14 @@ class RuntimeExpr(DelayedDeclerations):
         return self.__egg_typed_expr_thunk__()
 
     def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable | None:
-        cls_name = self.__egg_class_name__
-        class_decl = self.__egg_class_decl__
-
-        if name in (preserved_methods := class_decl.preserved_methods):
-            return preserved_methods[name].__get__(self)
-
-        if name in class_decl.methods:
-            return RuntimeFunction(self.__egg_decls_thunk__, Thunk.value(MethodRef(cls_name, name)), self)
-        if name in class_decl.properties:
-            return RuntimeFunction(self.__egg_decls_thunk__, Thunk.value(PropertyRef(cls_name, name)), self)()
-        raise AttributeError(f"{cls_name} has no method {name}") from None
+        if (method := _get_expr_method(self, name)) is not None:
+            return method
+        if name in self.__egg_class_decl__.properties:
+            fn = RuntimeFunction(
+                self.__egg_decls_thunk__, Thunk.value(PropertyRef(self.__egg_class_name__, name)), self
+            )
+            return fn()
+        raise AttributeError(f"{self.__egg_class_name__} has no method {name}") from None
 
     def __repr__(self) -> str:
         """
@@ -493,7 +541,7 @@ class RuntimeExpr(DelayedDeclerations):
         return pretty_decl(self.__egg_decls__, self.__egg_typed_expr__.expr, wrapping_fn=wrapping_fn)
 
     def _ipython_display_(self) -> None:
-        from IPython.display import Code, display
+        from IPython.display import Code, display  # noqa: PLC0415
 
         display(Code(str(self), language="python"))
 
@@ -509,13 +557,6 @@ class RuntimeExpr(DelayedDeclerations):
     def __egg_class_decl__(self) -> ClassDecl:
         return self.__egg_decls__.get_class_decl(self.__egg_class_name__)
 
-    # These both will be overriden below in the special methods section, but add these here for type hinting purposes
-    def __eq__(self, other: object) -> Fact:  # type: ignore[override, empty-body]
-        ...
-
-    def __ne__(self, other: object) -> RuntimeExpr:  # type: ignore[override, empty-body]
-        ...
-
     # Implement these so that copy() works on this object
     # otherwise copy will try to call `__getstate__` before object is initialized with properties which will cause inifinite recursion
 
@@ -527,91 +568,102 @@ class RuntimeExpr(DelayedDeclerations):
         self.__egg_typed_expr_thunk__ = Thunk.value(d[1])
 
     def __hash__(self) -> int:
+        if (method := _get_expr_method(self, "__hash__")) is not None:
+            return cast("int", cast("Any", method()))
         return hash(self.__egg_typed_expr__)
 
+    # Implement this directly to special case behavior where it transforms to an egraph equality, if it is not a
+    # preserved method or defined on the class
+    def __eq__(self, other: object) -> object:  # type: ignore[override]
+        if (method := _get_expr_method(self, "__eq__")) is not None:
+            return method(other)
 
-# Define each of the special methods, since we have already declared them for pretty printing
-for name in list(BINARY_METHODS) + list(UNARY_METHODS) + ["__getitem__", "__call__", "__setitem__", "__delitem__"]:
+        # TODO: Check if two objects can be upcasted to be the same. If not, then return NotImplemented so other
+        # expr gets a chance to resolve __eq__ which could be a preserved method.
+        from .egraph import BaseExpr, eq  # noqa: PLC0415
 
-    def _special_method(
-        self: RuntimeExpr,
-        *args: object,
-        __name: str = name,
-        **kwargs: object,
-    ) -> RuntimeExpr | Fact | None:
-        from .conversion import ConvertError
+        return eq(cast("BaseExpr", self)).to(cast("BaseExpr", other))
 
-        class_name = self.__egg_class_name__
-        class_decl = self.__egg_class_decl__
-        # First, try to resolve as preserved method
-        try:
-            method = class_decl.preserved_methods[__name]
-        except KeyError:
-            pass
-        else:
-            return method(self, *args, **kwargs)
-        # If this is a "partial" method meaning that it can return NotImplemented,
-        # we want to find the "best" superparent (lowest cost) of the arg types to call with it, instead of just
-        # using the arg type of the self arg.
-        # This is neccesary so if we add like an int to a ndarray, it will upcast the int to an ndarray, instead of vice versa.
-        if __name in PARTIAL_METHODS:
-            try:
-                return call_method_min_conversion(self, args[0], __name)
-            except ConvertError:
-                # Defer raising not imeplemented in case the dunder method is not symmetrical, then
-                # we use the standard process
-                pass
-        if __name in class_decl.methods:
-            fn = RuntimeFunction(self.__egg_decls_thunk__, Thunk.value(MethodRef(class_name, __name)), self)
-            return fn(*args, **kwargs)  # type: ignore[arg-type]
-        # Handle == and != fallbacks to eq and ne helpers if the methods aren't defined on the class explicitly.
-        if __name == "__eq__":
-            from .egraph import BaseExpr, eq
-
-            return eq(cast("BaseExpr", self)).to(cast("BaseExpr", args[0]))
-        if __name == "__ne__":
-            from .egraph import BaseExpr, ne
-
-            return cast("RuntimeExpr", ne(cast("BaseExpr", self)).to(cast("BaseExpr", args[0])))
-
-        if __name in PARTIAL_METHODS:
-            return NotImplemented
-        raise TypeError(f"{class_name!r} object does not support {__name}")
+    def __ne__(self, other: object) -> object:  # type: ignore[override]
+        if (method := _get_expr_method(self, "__ne__")) is not None:
+            return method(other)
+
+        from .egraph import BaseExpr, ne  # noqa: PLC0415
 
-    setattr(RuntimeExpr, name, _special_method)
+        return ne(cast("BaseExpr", self)).to(cast("BaseExpr", other))
 
-# For each of the reflected binary methods, translate to the corresponding non-reflected method
-for reflected, non_reflected in REFLECTED_BINARY_METHODS.items():
+    def __call__(
+        self, *args: object, **kwargs: object
+    ) -> object:  # define it here only for type checking, it will be overriden below
+        ...
 
-    def _reflected_method(self: RuntimeExpr, other: object, __non_reflected: str = non_reflected) -> RuntimeExpr | None:
-        # All binary methods are also "partial" meaning we should try to upcast first.
-        return call_method_min_conversion(other, self, __non_reflected)
 
-    setattr(RuntimeExpr, reflected, _reflected_method)
+def _get_expr_method(expr: RuntimeExpr, name: str) -> RuntimeFunction | RuntimeExpr | Callable | None:
+    if name in (preserved_methods := expr.__egg_class_decl__.preserved_methods):
+        return preserved_methods[name].__get__(expr)
 
+    if name in expr.__egg_class_decl__.methods:
+        return RuntimeFunction(expr.__egg_decls_thunk__, Thunk.value(MethodRef(expr.__egg_class_name__, name)), expr)
+    return None
 
-def call_method_min_conversion(slf: object, other: object, name: str) -> RuntimeExpr | None:
-    from .conversion import min_convertable_tp, resolve_literal
 
-    # find a minimum type that both can be converted to
-    # This is so so that calls like `-0.1 * Int("x")` work by upcasting both to floats.
-    min_tp = min_convertable_tp(slf, other, name).to_var()
-    slf = resolve_literal(min_tp, slf)
-    other = resolve_literal(min_tp, other)
-    method = RuntimeFunction(slf.__egg_decls_thunk__, Thunk.value(MethodRef(slf.__egg_class_name__, name)), slf)
-    return method(other)
+def define_expr_method(name: str) -> None:
+    """
+    Given the name of a method, explicitly defines it on the runtime type that holds `Expr` objects as a method.
 
+    Call this if you need a method to be defined on the type itself where overrindg with `__getattr__` does not suffice,
+    like for NumPy's `__array_ufunc__`.
+    """
 
-for name in PRESERVED_METHODS:
+    def _defined_method(self: RuntimeExpr, *args, __name: str = name, **kwargs):
+        fn = _get_expr_method(self, __name)
+        if fn is None:
+            raise TypeError(f"{self.__egg_class_name__} expression has no method {__name}")
+        return fn(*args, **kwargs)
 
-    def _preserved_method(self: RuntimeExpr, __name: str = name):
-        try:
-            method = self.__egg_decls__.get_class_decl(self.__egg_typed_expr__.tp.name).preserved_methods[__name]
-        except KeyError as e:
-            raise TypeError(f"{self.__egg_typed_expr__.tp.name} has no method {__name}") from e
-        return method(self)
+    setattr(RuntimeExpr, name, _defined_method)
+
+
+for name in TYPE_DEFINED_METHODS:
+    define_expr_method(name)
+
+
+for name, r_method in itertools.product(NUMERIC_BINARY_METHODS, (False, True)):
+
+    def _numeric_binary_method(self: object, other: object, name: str = name, r_method: bool = r_method) -> object:
+        """
+        Implements numeric binary operations.
+
+        Tries to find the minimum cost conversion of either the LHS or the RHS, by finding all methods with either
+        the LHS or the RHS as exactly the right type and then upcasting the other to that type.
+        """
+        # 1. switch if reversed method
+        if r_method:
+            self, other = other, self
+        # If the types don't exactly match to start, then we need to try converting one of them, by finding the cheapest conversion
+        if not (
+            isinstance(self, RuntimeExpr)
+            and isinstance(other, RuntimeExpr)
+            and (
+                self.__egg_decls__.check_binary_method_with_types(
+                    name, self.__egg_typed_expr__.tp, other.__egg_typed_expr__.tp
+                )
+            )
+        ):
+            from .conversion import min_binary_conversion, resolve_type  # noqa: PLC0415
+
+            best_method = min_binary_conversion(name, resolve_type(self), resolve_type(other))
+
+            if not best_method:
+                raise RuntimeError(f"Cannot resolve {name} for {self} and {other}, no conversion found")
+            self, other = best_method[0](self), best_method[1](other)
+
+        method_ref = MethodRef(self.__egg_class_name__, name)
+        fn = RuntimeFunction(Thunk.value(self.__egg_decls__), Thunk.value(method_ref), self)
+        return fn(other)
 
-    setattr(RuntimeExpr, name, _preserved_method)
+    method_name = f"__r{name[2:]}" if r_method else name
+    setattr(RuntimeExpr, method_name, _numeric_binary_method)
 
 
 def resolve_callable(callable: object) -> tuple[CallableRef, Declarations]:

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 @@ class Thunk(Generic[T, Unpack[TS]]):
         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

egglog/type_constraint_solver.py

@@ -107,7 +107,7 @@ class TypeConstraintSolver:
                 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)

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)):  # 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  # 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]

{egglog-10.0.1.dist-info → egglog-11.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: egglog
-Version: 10.0.1
+Version: 11.0.0
 Classifier: Environment :: MacOS X
 Classifier: Environment :: Win32 (MS Windows)
 Classifier: Intended Audience :: Developers