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/changelog.md

@@ -9,6 +9,8 @@ _This project uses semantic versioning_
     Also changes the representation to be an index into a list instead of the ID, making egglog programs more deterministic.
   - Prefix constant declerations and unbound variables to not shadow let variables
   - BREAKING: Remove `simplify` since it was removed upstream. You can manually replace it with an insert, run, then extract.
+- Change how anonymous functions are converted to remove metaprogramming and lift only the unbound variables as args
+- Add support for getting the "value" of a function type with `.eval()`, i.e. `assert UnstableFn(f).eval() == f`.
 
 ## 10.0.2 (2025-06-22)
 

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/builtins.py

@@ -8,16 +8,18 @@
 from collections.abc import Callable
 from fractions import Fraction
 from functools import partial, reduce
+from inspect import signature
 from types import FunctionType, MethodType
 from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, cast, overload
 
 from typing_extensions import TypeVarTuple, Unpack
 
-from .conversion import convert, converter, get_type_args
+from egglog.declarations import TypedExprDecl
+
+from .conversion import convert, converter, get_type_args, resolve_literal
 from .declarations import *
-from .egraph import BaseExpr, BuiltinExpr, expr_fact, function, get_current_ruleset, method
-from .functionalize import functionalize
-from .runtime import RuntimeClass, RuntimeExpr, RuntimeFunction
+from .egraph import BaseExpr, BuiltinExpr, _add_default_rewrite_inner, expr_fact, function, get_current_ruleset, method
+from .runtime import RuntimeClass, RuntimeExpr, RuntimeFunction, resolve_type_annotation_mutate
 from .thunk import Thunk
 
 if TYPE_CHECKING:
@@ -1022,46 +1024,89 @@ def __init__(self, f, *partial) -> None: ...
     @method(egg_fn="unstable-app")
     def __call__(self, *args: Unpack[TS]) -> T: ...
 
+    @method(preserve=True)
+    def eval(self) -> Callable[[Unpack[TS]], T]:
+        """
+        If this is a constructor, returns either the callable directly or a `functools.partial` function if args are provided.
+        """
+        assert isinstance(self, RuntimeExpr)
+        match self.__egg_typed_expr__.expr:
+            case PartialCallDecl(CallDecl() as call):
+                fn, args = _deconstruct_call_decl(self.__egg_decls_thunk__, call)
+                if not args:
+                    return fn
+                return partial(fn, *args)
+        msg = "UnstableFn can only be evaluated if it is a function or a partial application of a function."
+        raise BuiltinEvalError(msg)
+
+
+def _deconstruct_call_decl(
+    decls_thunk: Callable[[], Declarations], call: CallDecl
+) -> tuple[Callable, tuple[object, ...]]:
+    """
+    Deconstructs a CallDecl into a runtime callable and its arguments.
+    """
+    args = call.args
+    arg_exprs = tuple(RuntimeExpr(decls_thunk, Thunk.value(a)) for a in args)
+    egg_bound = (
+        JustTypeRef(call.callable.class_name, call.bound_tp_params or ())
+        if isinstance(call.callable, (ClassMethodRef, InitRef, ClassVariableRef))
+        else None
+    )
+    if isinstance(call.callable, InitRef):
+        return RuntimeClass(
+            decls_thunk,
+            TypeRefWithVars(
+                call.callable.class_name,
+            ),
+        ), arg_exprs
+    return RuntimeFunction(decls_thunk, Thunk.value(call.callable), egg_bound), arg_exprs
+
 
 # Method Type is for builtins like __getitem__
 converter(MethodType, UnstableFn, lambda m: UnstableFn(m.__func__, m.__self__))
 converter(RuntimeFunction, UnstableFn, UnstableFn)
 converter(partial, UnstableFn, lambda p: UnstableFn(p.func, *p.args))
 
 
-def _convert_function(a: FunctionType) -> UnstableFn:
+def _convert_function(fn: FunctionType) -> UnstableFn:
     """
-    Converts a function type to an unstable function
+    Converts a function type to an unstable function. This function will be an anon function in egglog.
 
-    Would just be UnstableFn(function(a)) but we have to look for any nonlocals and globals
-    which are runtime expressions with `var`s in them and add them as args to the function
+    Would just be UnstableFn(function(a)) but we have to account for unbound vars within the body.
+
+    This means that we have to turn all of those unbound vars into args to the function, and then
+    partially apply them, alongside creating a default rewrite for the function.
     """
-    # Update annotations of a to be the type we are trying to convert to
-    return_tp, *arg_tps = get_type_args()
-    a.__annotations__ = {
-        "return": return_tp,
-        # The first varnames should always be the arg names
-        **dict(zip(a.__code__.co_varnames, arg_tps, strict=False)),
-    }
-    # Modify name to make it unique
-    # a.__name__ = f"{a.__name__} {hash(a.__code__)}"
-    transformed_fn = functionalize(a, value_to_annotation)
-    assert isinstance(transformed_fn, partial)
-    return UnstableFn(
-        function(ruleset=get_current_ruleset(), use_body_as_name=True, subsume=True)(transformed_fn.func),
-        *transformed_fn.args,
+    decls = Declarations()
+    return_type, *arg_types = [resolve_type_annotation_mutate(decls, tp) for tp in get_type_args()]
+    arg_names = [p.name for p in signature(fn).parameters.values()]
+    arg_decls = [
+        TypedExprDecl(tp.to_just(), UnboundVarDecl(name)) for name, tp in zip(arg_names, arg_types, strict=True)
+    ]
+    res = resolve_literal(
+        return_type, fn(*(RuntimeExpr.__from_values__(decls, a) for a in arg_decls)), Thunk.value(decls)
     )
+    res_expr = res.__egg_typed_expr__
+    decls |= res
+    # these are all the args that appear in the body that are not bound by the args of the function
+    unbound_vars = list(collect_unbound_vars(res_expr) - set(arg_decls))
+    # prefix the args with them
+    fn_ref = UnnamedFunctionRef(tuple(unbound_vars + arg_decls), res_expr)
+    rewrite_decl = DefaultRewriteDecl(fn_ref, res_expr.expr, subsume=True)
+    ruleset_decls = _add_default_rewrite_inner(decls, rewrite_decl, get_current_ruleset())
+    ruleset_decls |= res
 
-
-def value_to_annotation(a: object) -> type | None:
-    # only lift runtime expressions (which could contain vars) not any other nonlocals/globals we use in the function
-    if not isinstance(a, RuntimeExpr):
-        return None
-    return cast("type", RuntimeClass(Thunk.value(a.__egg_decls__), a.__egg_typed_expr__.tp.to_var()))
+    fn = RuntimeFunction(Thunk.value(decls), Thunk.value(fn_ref))
+    return UnstableFn(fn, *(RuntimeExpr.__from_values__(decls, v) for v in unbound_vars))
 
 
 converter(FunctionType, UnstableFn, _convert_function)
 
+##
+# Utility Functions
+##
+
 
 def _extract_lit(e: BaseExpr) -> LitType:
     """

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)