Add support for adding methods actually defined on classes

Adds support for methods like `__array_function__` which actually need to be added on the class as actual methods, not through overloading `__getattr__`. Custom methods can be registered by third party libraries.

This PR also redoes the logic for upcasting when using binary operations. Instead of upcasting both values, it will only ever upcast one, choosing whichever one would be cheaper to upcast. This leads to more predictable behavior.

Author: saulshanabrook

docs/reference/python-integration.md

@@ -303,6 +303,11 @@ Note that the following list of methods are only supported as "preserved" since
 - `__iter_`
 - `__index__`
 
+If you want to register additional methods as always preserved and defined on the `Expr` class itself, if needed
+instead of the normal mechanism which relies on `__getattr__`, you can call `egglog.define_expr_method(name: str)`,
+with the name of a method. This is only needed for third party code that inspects the type object itself to see if a
+method is defined instead of just attempting to call it.
+
 ### Reflected methods
 
 Note that reflected methods (i.e. `__radd__`) are handled as a special case. If defined, they won't create their own egglog functions.

python/egglog/__init__.py

@@ -5,7 +5,8 @@
 from . import config, ipython_magic  # noqa: F401
 from .bindings import EggSmolError  # noqa: F401
 from .builtins import *  # noqa: UP029
-from .conversion import ConvertError, convert, converter, get_type_args  # noqa: F401
+from .conversion import *
 from .egraph import *
+from .runtime import define_expr_method as define_expr_method  # noqa: PLC0414
 
 del ipython_magic

python/egglog/conversion.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 from collections import defaultdict
+from collections.abc import Callable
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, TypeVar, cast
+from typing import TYPE_CHECKING, Any, TypeVar, cast
 
 from .declarations import *
 from .pretty import *
@@ -13,22 +14,22 @@
 from .type_constraint_solver import TypeConstraintError
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Generator
+    from collections.abc import Generator
 
     from .egraph import BaseExpr
     from .type_constraint_solver import TypeConstraintSolver
 
-__all__ = ["ConvertError", "convert", "convert_to_same_type", "converter", "resolve_literal"]
+__all__ = ["ConvertError", "convert", "converter", "get_type_args"]
 # Mapping from (source type, target type) to and function which takes in the runtimes values of the source and return the target
-CONVERSIONS: dict[tuple[type | JustTypeRef, JustTypeRef], tuple[int, Callable]] = {}
+CONVERSIONS: dict[tuple[type | JustTypeRef, JustTypeRef], tuple[int, Callable[[Any], RuntimeExpr]]] = {}
 # Global declerations to store all convertable types so we can query if they have certain methods or not
 _CONVERSION_DECLS = Declarations.create()
 # Defer a list of declerations to be added to the global declerations, so that we can not trigger them procesing
 # until we need them
 _TO_PROCESS_DECLS: list[DeclerationsLike] = []
 
 
-def _retrieve_conversion_decls() -> Declarations:
+def retrieve_conversion_decls() -> Declarations:
     _CONVERSION_DECLS.update(*_TO_PROCESS_DECLS)
     _TO_PROCESS_DECLS.clear()
     return _CONVERSION_DECLS
@@ -49,10 +50,10 @@ def converter(from_type: type[T], to_type: type[V], fn: Callable[[T], V], cost:
     to_type_name = process_tp(to_type)
     if not isinstance(to_type_name, JustTypeRef):
         raise TypeError(f"Expected return type to be a egglog type, got {to_type_name}")
-    _register_converter(process_tp(from_type), to_type_name, fn, cost)
+    _register_converter(process_tp(from_type), to_type_name, cast("Callable[[Any], RuntimeExpr]", fn), cost)
 
 
-def _register_converter(a: type | JustTypeRef, b: JustTypeRef, a_b: Callable, cost: int) -> None:
+def _register_converter(a: type | JustTypeRef, b: JustTypeRef, a_b: Callable[[Any], RuntimeExpr], cost: int) -> None:
     """
     Registers a converter from some type to an egglog type, if not already registered.
 
@@ -97,15 +98,15 @@ class _ComposedConverter:
     We use the dataclass instead of the lambda to make it easier to debug.
     """
 
-    a_b: Callable
-    b_c: Callable
+    a_b: Callable[[Any], RuntimeExpr]
+    b_c: Callable[[Any], RuntimeExpr]
     b_args: tuple[JustTypeRef, ...]
 
-    def __call__(self, x: object) -> object:
+    def __call__(self, x: Any) -> RuntimeExpr:
         # if we have A -> B and B[C] -> D then we should use (C,) as the type args
         # when converting from A -> B
         if self.b_args:
-            with with_type_args(self.b_args, _retrieve_conversion_decls):
+            with with_type_args(self.b_args, retrieve_conversion_decls):
                 first_res = self.a_b(x)
         else:
             first_res = self.a_b(x)
@@ -142,33 +143,38 @@ def process_tp(tp: type | RuntimeClass) -> JustTypeRef | type:
     return tp
 
 
-def min_convertable_tp(a: object, b: object, name: str) -> JustTypeRef:
-    """
-    Returns the minimum convertable type between a and b, that has a method `name`, raising a ConvertError if no such type exists.
-    """
-    decls = _retrieve_conversion_decls()
-    a_tp = _get_tp(a)
-    b_tp = _get_tp(b)
-    # Make sure at least one of the types has the method, to avoid issue with == upcasting improperly
-    if not (
-        (isinstance(a_tp, JustTypeRef) and decls.has_method(a_tp.name, name))
-        or (isinstance(b_tp, JustTypeRef) and decls.has_method(b_tp.name, name))
-    ):
-        raise ConvertError(f"Neither {a_tp} nor {b_tp} has method {name}")
-    a_converts_to = {
-        to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == a_tp and decls.has_method(to.name, name)
-    }
-    b_converts_to = {
-        to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == b_tp and decls.has_method(to.name, name)
-    }
-    if isinstance(a_tp, JustTypeRef) and decls.has_method(a_tp.name, name):
-        a_converts_to[a_tp] = 0
-    if isinstance(b_tp, JustTypeRef) and decls.has_method(b_tp.name, name):
-        b_converts_to[b_tp] = 0
-    common = set(a_converts_to) & set(b_converts_to)
-    if not common:
-        raise ConvertError(f"Cannot convert {a_tp} and {b_tp} to a common type")
-    return min(common, key=lambda tp: a_converts_to[tp] + b_converts_to[tp])
+# def min_convertable_tp(a: object, b: object, name: str) -> JustTypeRef:
+#     """
+#     Returns the minimum convertable type between a and b, that has a method `name`, raising a ConvertError if no such type exists.
+#     """
+#     decls = _retrieve_conversion_decls().copy()
+#     if isinstance(a, RuntimeExpr):
+#         decls |= a
+#     if isinstance(b, RuntimeExpr):
+#         decls |= b
+
+#     a_tp = _get_tp(a)
+#     b_tp = _get_tp(b)
+#     # Make sure at least one of the types has the method, to avoid issue with == upcasting improperly
+#     if not (
+#         (isinstance(a_tp, JustTypeRef) and decls.has_method(a_tp.name, name))
+#         or (isinstance(b_tp, JustTypeRef) and decls.has_method(b_tp.name, name))
+#     ):
+#         raise ConvertError(f"Neither {a_tp} nor {b_tp} has method {name}")
+#     a_converts_to = {
+#         to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == a_tp and decls.has_method(to.name, name)
+#     }
+#     b_converts_to = {
+#         to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == b_tp and decls.has_method(to.name, name)
+#     }
+#     if isinstance(a_tp, JustTypeRef) and decls.has_method(a_tp.name, name):
+#         a_converts_to[a_tp] = 0
+#     if isinstance(b_tp, JustTypeRef) and decls.has_method(b_tp.name, name):
+#         b_converts_to[b_tp] = 0
+#     common = set(a_converts_to) & set(b_converts_to)
+#     if not common:
+#         raise ConvertError(f"Cannot convert {a_tp} and {b_tp} to a common type")
+#     return min(common, key=lambda tp: a_converts_to[tp] + b_converts_to[tp])
 
 
 def identity(x: object) -> object:
@@ -197,7 +203,7 @@ def with_type_args(args: tuple[JustTypeRef, ...], decls: Callable[[], Declaratio
 def resolve_literal(
     tp: TypeOrVarRef,
     arg: object,
-    decls: Callable[[], Declarations] = _retrieve_conversion_decls,
+    decls: Callable[[], Declarations] = retrieve_conversion_decls,
     tcs: TypeConstraintSolver | None = None,
     cls_name: str | None = None,
 ) -> RuntimeExpr:
@@ -208,12 +214,12 @@ def resolve_literal(
 
     If it cannot be resolved, we assume that the value passed in will resolve it.
     """
-    arg_type = _get_tp(arg)
+    arg_type = resolve_type(arg)
 
     # If we have any type variables, dont bother trying to resolve the literal, just return the arg
     try:
         tp_just = tp.to_just()
-    except NotImplementedError:
+    except TypeVarError:
         # If this is a generic arg but passed in a non runtime expression, try to resolve the generic
         # args first based on the existing type constraint solver
         if tcs:
@@ -258,7 +264,7 @@ def _debug_print_converers():
         source_to_targets[source].append(target)
 
 
-def _get_tp(x: object) -> JustTypeRef | type:
+def resolve_type(x: object) -> JustTypeRef | type:
     if isinstance(x, RuntimeExpr):
         return x.__egg_typed_expr__.tp
     tp = type(x)

python/egglog/declarations.py

@@ -73,6 +73,7 @@
     "SpecialFunctions",
     "TypeOrVarRef",
     "TypeRefWithVars",
+    "TypeVarError",
     "TypedExprDecl",
     "UnboundVarDecl",
     "UnionDecl",
@@ -95,7 +96,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 declarations for {self}"
+            msg = f"Cannot resolve declarations {err}"
             raise RuntimeError(msg) from err
 
 
@@ -225,14 +226,43 @@ def set_function_decl(
             case _:
                 assert_never(ref)
 
-    def has_method(self, class_name: str, method_name: str) -> bool | None:
+    def check_binary_method_with_types(self, method_name: str, self_type: JustTypeRef, other_type: JustTypeRef) -> bool:
         """
-        Returns whether the given class has the given method, or None if we cant find the class.
+        Checks if the class has a binary method compatible with the given types.
         """
-        if class_name in self._classes:
-            return method_name in self._classes[class_name].methods
+        vars: dict[ClassTypeVarRef, JustTypeRef] = {}
+        if callable_decl := self._classes[self_type.name].methods.get(method_name):
+            match callable_decl.signature:
+                case FunctionSignature((self_arg_type, other_arg_type)) if self_arg_type.matches_just(
+                    vars, self_type
+                ) and other_arg_type.matches_just(vars, other_type):
+                    return True
+        return False
+
+    def check_binary_method_with_self_type(self, method_name: str, self_type: JustTypeRef) -> JustTypeRef | None:
+        """
+        Checks if the class has a binary method with the given name and self type. Returns the other type if it exists.
+        """
+        vars: dict[ClassTypeVarRef, JustTypeRef] = {}
+        if callable_decl := self._classes[self_type.name].methods.get(method_name):
+            match callable_decl.signature:
+                case FunctionSignature((self_arg_type, other_arg_type)) if self_arg_type.matches_just(vars, self_type):
+                    return other_arg_type.to_just(vars)
         return None
 
+    def check_binary_method_with_other_type(self, method_name: str, other_type: JustTypeRef) -> Iterable[JustTypeRef]:
+        """
+        Returns the types which are compatible with the given binary method name and other type.
+        """
+        for class_decl in self._classes.values():
+            vars: dict[ClassTypeVarRef, JustTypeRef] = {}
+            if callable_decl := class_decl.methods.get(method_name):
+                match callable_decl.signature:
+                    case FunctionSignature((self_arg_type, other_arg_type)) if other_arg_type.matches_just(
+                        vars, other_type
+                    ):
+                        yield self_arg_type.to_just(vars)
+
     def get_class_decl(self, name: str) -> ClassDecl:
         return self._classes[name]
 
@@ -300,6 +330,10 @@ def __str__(self) -> str:
 _RESOLVED_TYPEVARS: dict[ClassTypeVarRef, TypeVar] = {}
 
 
+class TypeVarError(RuntimeError):
+    """Error when trying to resolve a type variable that doesn't exist."""
+
+
 @dataclass(frozen=True)
 class ClassTypeVarRef:
     """
@@ -309,9 +343,10 @@ class ClassTypeVarRef:
     name: str
     module: str
 
-    def to_just(self) -> JustTypeRef:
-        msg = f"{self}: egglog does not support generic classes yet."
-        raise NotImplementedError(msg)
+    def to_just(self, vars: dict[ClassTypeVarRef, JustTypeRef] | None = None) -> JustTypeRef:
+        if vars is None or self not in vars:
+            raise TypeVarError(f"Cannot convert {self} to just type")
+        return vars[self]
 
     def __str__(self) -> str:
         return str(self.to_type_var())
@@ -325,20 +360,39 @@ def from_type_var(cls, typevar: TypeVar) -> ClassTypeVarRef:
     def to_type_var(self) -> TypeVar:
         return _RESOLVED_TYPEVARS[self]
 
+    def matches_just(self, vars: dict[ClassTypeVarRef, JustTypeRef], other: JustTypeRef) -> bool:
+        """
+        Checks if this type variable matches the given JustTypeRef, including type variables.
+        """
+        if self in vars:
+            return vars[self] == other
+        vars[self] = other
+        return True
+
 
 @dataclass(frozen=True)
 class TypeRefWithVars:
     name: str
     args: tuple[TypeOrVarRef, ...] = ()
 
-    def to_just(self) -> JustTypeRef:
-        return JustTypeRef(self.name, tuple(a.to_just() for a in self.args))
+    def to_just(self, vars: dict[ClassTypeVarRef, JustTypeRef] | None = None) -> JustTypeRef:
+        return JustTypeRef(self.name, tuple(a.to_just(vars) for a in self.args))
 
     def __str__(self) -> str:
         if self.args:
             return f"{self.name}[{', '.join(str(a) for a in self.args)}]"
         return self.name
 
+    def matches_just(self, vars: dict[ClassTypeVarRef, JustTypeRef], other: JustTypeRef) -> bool:
+        """
+        Checks if this type reference matches the given JustTypeRef, including type variables.
+        """
+        return (
+            self.name == other.name
+            and len(self.args) == len(other.args)
+            and all(a.matches_just(vars, b) for a, b in zip(self.args, other.args, strict=True))
+        )
+
 
 TypeOrVarRef: TypeAlias = ClassTypeVarRef | TypeRefWithVars
 

python/egglog/egraph.py

@@ -115,7 +115,7 @@
     "__firstlineno__",
     "__static_attributes__",
     # Ignore all reflected binary method
-    *REFLECTED_BINARY_METHODS.keys(),
+    *(f"__r{m[2:]}" for m in NUMERIC_BINARY_METHODS),
 }