egglog 12.0.0__cp313-cp313t-manylinux_2_17_ppc64.manylinux2014_ppc64.whl

egglog/runtime.py

@@ -0,0 +1,888 @@
+"""
+Holds a number of types which are only used at runtime to emulate Python objects.
+
+Users will not import anything from this module, and statically they won't know these are the types they are using.
+
+But at runtime they will be exposed.
+
+Note that all their internal fields are prefixed with __egg_ to avoid name collisions with user code, but will end in __
+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 InitVar, dataclass, replace
+from inspect import Parameter, Signature
+from itertools import zip_longest
+from typing import TYPE_CHECKING, Any, TypeVar, Union, assert_never, cast, get_args, get_origin
+
+import cloudpickle
+
+from .declarations import *
+from .pretty import *
+from .thunk import Thunk
+from .type_constraint_solver import *
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+
+__all__ = [
+    "ALWAYS_MUTATES_SELF",
+    "ALWAYS_PRESERVED",
+    "LIT_IDENTS",
+    "NUMERIC_BINARY_METHODS",
+    "RuntimeClass",
+    "RuntimeExpr",
+    "RuntimeFunction",
+    "create_callable",
+    "define_expr_method",
+    "resolve_callable",
+    "resolve_type_annotation",
+    "resolve_type_annotation_mutate",
+]
+
+
+UNIT_IDENT = Ident.builtin("Unit")
+UNARY_LIT_IDENTS = {Ident.builtin("i64"), Ident.builtin("f64"), Ident.builtin("Bool"), Ident.builtin("String")}
+LIT_IDENTS = UNARY_LIT_IDENTS | {UNIT_IDENT, Ident.builtin("PyObject")}
+
+# 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__",
+}
+
+
+# special methods that return none and mutate self
+ALWAYS_MUTATES_SELF = {
+    "__setitem__",
+    "__delitem__",
+    # "__setattr__",
+}
+
+# special methods which must return real python values instead of lazy expressions
+ALWAYS_PRESERVED = {
+    "__bytes__",
+    # "__format__",
+    "__hash__",
+    "__bool__",
+    "__len__",
+    "__length_hint__",
+    "__iter__",
+    "__reversed__",
+    "__contains__",
+    "__index__",
+    "__buffer__",
+    "__complex__",
+    "__int__",
+    "__float__",
+    "__str__",
+    "__repr__",
+}
+
+# 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 = {
+    "__call__",
+    "__getitem__",
+    "__pos__",
+    "__neg__",
+    "__invert__",
+    "__round__",
+    "__abs__",
+    "__trunc__",
+    "__floor__",
+    "__ceil__",
+    *ALWAYS_PRESERVED,
+    *ALWAYS_MUTATES_SELF,
+} - {"__str__", "__repr__"}  # these are defined on RuntimeFunction
+
+# Methods that need to be defined as descriptors on the class so that they can be looked up on the class itself.
+CLASS_DESCRIPTOR_METHODS = {
+    "__eq__",
+    "__ne__",
+    "__str__",
+    "__repr__",
+    "__getattr__",
+    *NUMERIC_BINARY_METHODS,
+    *TYPE_DEFINED_METHODS,
+}
+
+# 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.
+_PY_OBJECT_CLASS: RuntimeClass | None = None
+# Same for functions
+_UNSTABLE_FN_CLASS: RuntimeClass | None = None
+
+T = TypeVar("T")
+
+
+def resolve_type_annotation_mutate(decls: Declarations, tp: object) -> TypeOrVarRef:
+    """
+    Wrap resolve_type_annotation to mutate decls, as a helper for internal use in sitations where that is more ergonomic.
+    """
+    new_decls, tp = resolve_type_annotation(tp)
+    decls |= new_decls
+    return tp
+
+
+def resolve_type_annotation(tp: object) -> tuple[DeclerationsLike, TypeOrVarRef]:
+    """
+    Resolves a type object into a type reference.
+
+    Any runtime type object decls will be returned as well. We do this so we can use this without having to
+    resolve the decls if need be.
+    """
+    if isinstance(tp, TypeVar):
+        return None, ClassTypeVarRef.from_type_var(tp)
+    # If there is a union, then we assume the first item is the type we want, and the others are types that can be converted to that type.
+    if get_origin(tp) == Union:
+        first, *_rest = get_args(tp)
+        return resolve_type_annotation(first)
+
+    # If the type is `object` then this is assumed to be a PyObjectLike, i.e. converted into a PyObject
+    if tp is object:
+        assert _PY_OBJECT_CLASS
+        return resolve_type_annotation(_PY_OBJECT_CLASS)
+    # If the type is a `Callable` then convert it into a UnstableFn
+    if get_origin(tp) == Callable:
+        assert _UNSTABLE_FN_CLASS
+        args, ret = get_args(tp)
+        return resolve_type_annotation(_UNSTABLE_FN_CLASS[(ret, *args)])
+    if isinstance(tp, RuntimeClass):
+        return tp, tp.__egg_tp__
+    raise TypeError(f"Unexpected type annotation {tp}")
+
+
+def inverse_resolve_type_annotation(decls_thunk: Callable[[], Declarations], tp: TypeOrVarRef) -> object:
+    """
+    Inverse of resolve_type_annotation
+    """
+    if isinstance(tp, ClassTypeVarRef):
+        return tp.to_type_var()
+    return RuntimeClass(decls_thunk, tp)
+
+
+##
+# Runtime objects
+##
+
+
+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__.ident == instance.__egg_typed_expr__.tp.ident
+
+
+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__):
+            if m is not object:
+                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}, lambda ns: ns.update(RUNTIME_CLASS_DESCRIPTORS))
+        init(tp, *args, **kwargs)
+        return tp
+
+    def __instancecheck__(cls, instance: object) -> bool:
+        return isinstance(instance, BaseClassFactoryMeta)
+
+
+# Add a custom descriptor for each method we want to support on the class itself, so that we can lookup things like __add__
+# on the expr class. This is used for things like running docstrings.
+
+
+@dataclass
+class RuntimeClassDescriptor:
+    name: str
+
+    def __get__(self, obj: object, owner: RuntimeClass | None = None) -> Callable:
+        if owner is None:
+            raise AttributeError(f"Can only access {self.name} on the class, not an instance")
+        cls_decl = owner.__egg_decls__._classes[owner.__egg_tp__.ident]
+        if self.name in cls_decl.class_methods:
+            return RuntimeFunction(
+                owner.__egg_decls_thunk__, Thunk.value(ClassMethodRef(owner.__egg_tp__.ident, self.name)), None
+            )
+        if self.name in cls_decl.preserved_methods:
+            return cls_decl.preserved_methods[self.name]
+        raise AttributeError(f"Class {owner.__egg_tp__.ident} has no method {self.name}") from None
+
+
+RUNTIME_CLASS_DESCRIPTORS: dict[str, RuntimeClassDescriptor] = {
+    method: RuntimeClassDescriptor(method) for method in CLASS_DESCRIPTOR_METHODS
+}
+
+
+@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, _egg_has_params: bool) -> None:
+        global _PY_OBJECT_CLASS, _UNSTABLE_FN_CLASS
+        ident = self.__egg_tp__.ident
+        if (ident) == Ident.builtin("PyObject"):
+            _PY_OBJECT_CLASS = self
+        elif ident == Ident.builtin("UnstableFn") and not self.__egg_tp__.args:
+            _UNSTABLE_FN_CLASS = self
+        # Set the module and name so that this works with doctest
+        if ident.module:
+            self.__module__ = ident.module
+        self.__name__ = ident.name
+
+    def verify(self) -> None:
+        if not self.__egg_tp__.args:
+            return
+
+        # Raise error if we have args, but they are the wrong number
+        desired_args = self.__egg_decls__.get_class_decl(self.__egg_tp__.ident).type_vars
+        if len(self.__egg_tp__.args) != len(desired_args):
+            raise ValueError(f"Expected {desired_args} type args, got {len(self.__egg_tp__.args)}")
+
+    def __call__(self, *args: object, **kwargs: object) -> RuntimeExpr | None:
+        """
+        Create an instance of this kind by calling the __init__ classmethod
+        """
+        # If this is a literal type, initializing it with a literal should return a literal
+        if (name := self.__egg_tp__.ident) == Ident.builtin("PyObject"):
+            assert len(args) == 1
+            try:
+                pickled = cloudpickle.dumps(args[0])
+            except Exception as e:
+                e.add_note(f"Failed to pickle object of type {type(args[0])}")
+                raise
+            return RuntimeExpr(
+                self.__egg_decls_thunk__, Thunk.value(TypedExprDecl(self.__egg_tp__.to_just(), PyObjectDecl(pickled)))
+            )
+        if name == Ident.builtin("UnstableFn"):
+            assert not kwargs
+            fn_arg, *partial_args = args
+            del args
+            # Assumes we don't have types set for UnstableFn w/ generics, that they have to be inferred
+
+            # 1. Call it with the partial args, and use untyped vars for the rest of the args
+            res = cast("Callable", fn_arg)(*partial_args, _egg_partial_function=True)
+            assert res is not None, "Mutable partial functions not supported"
+            # 2. Use the inferred return type and inferred rest arg types as the types of the function, and
+            #    the partially applied args as the args.
+            call = (res_typed_expr := res.__egg_typed_expr__).expr
+            return_tp = res_typed_expr.tp
+            assert isinstance(call, CallDecl), "partial function must be a call"
+            n_args = len(partial_args)
+            value = PartialCallDecl(replace(call, args=call.args[:n_args]))
+            remaining_arg_types = [a.tp for a in call.args[n_args:]]
+            type_ref = JustTypeRef(Ident.builtin("UnstableFn"), (return_tp, *remaining_arg_types))
+            return RuntimeExpr.__from_values__(Declarations.create(self, res), TypedExprDecl(type_ref, value))
+
+        if name in UNARY_LIT_IDENTS:
+            assert len(args) == 1
+            assert isinstance(args[0], int | float | str | bool)
+            return RuntimeExpr(
+                self.__egg_decls_thunk__, Thunk.value(TypedExprDecl(self.__egg_tp__.to_just(), LitDecl(args[0])))
+            )
+        if name == UNIT_IDENT:
+            assert len(args) == 0
+            return RuntimeExpr(
+                self.__egg_decls_thunk__, Thunk.value(TypedExprDecl(self.__egg_tp__.to_just(), LitDecl(None)))
+            )
+        fn = RuntimeFunction(self.__egg_decls_thunk__, Thunk.value(InitRef(name)), self.__egg_tp__.to_just())
+        return fn(*args, **kwargs)  # type: ignore[arg-type]
+
+    def __dir__(self) -> Iterable[str]:
+        cls_decl = self.__egg_decls__._classes[self.__egg_tp__.ident]
+        return (
+            list(cls_decl.class_methods)
+            + list(cls_decl.class_variables)
+            + list(cls_decl.preserved_methods)
+            + list(cls_decl.methods)
+            + list(cls_decl.properties)
+        )
+
+    def __getitem__(self, args: object) -> RuntimeClass:
+        if not isinstance(args, tuple):
+            args = (args,)
+        # defer resolving decls so that we can do generic instantiation for converters before all
+        # method types are defined.
+        decls_like, new_args = cast(
+            "tuple[tuple[DeclerationsLike, ...], tuple[TypeOrVarRef, ...]]",
+            zip(*(resolve_type_annotation(arg) for arg in args), strict=False),
+        )
+        # if we already have some args bound and some not, then we shold replace all existing args of typevars with new
+        # args
+        if old_args := self.__egg_tp__.args:
+            is_typevar = [isinstance(arg, ClassTypeVarRef) for arg in old_args]
+            if sum(is_typevar) != len(new_args):
+                raise TypeError(f"Expected {sum(is_typevar)} typevars, got {len(new_args)}")
+            new_args_list = list(new_args)
+            final_args = tuple(new_args_list.pop(0) if is_typevar[i] else old_args[i] for i in range(len(old_args)))
+        else:
+            final_args = new_args
+        tp = TypeRefWithVars(self.__egg_tp__.ident, final_args)
+        return RuntimeClass(Thunk.fn(Declarations.create, self, *decls_like), tp, _egg_has_params=True)
+
+    def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable:
+        if not isinstance(name, str):
+            raise TypeError(f"Attribute name must be a string, got {name!r}")
+        if name == "__origin__" and self.__egg_tp__.args:
+            return RuntimeClass(self.__egg_decls_thunk__, TypeRefWithVars(self.__egg_tp__.ident))
+
+        # Special case some names that don't exist so we can exit early without resolving decls
+        # Important so if we take union of RuntimeClass it won't try to resolve decls
+        if name in {
+            "__typing_subst__",
+            "__parameters__",
+            # Origin is used in get_type_hints which is used when resolving the class itself
+            "__origin__",
+            "__typing_unpacked_tuple_args__",
+            "__typing_is_unpacked_typevartuple__",
+        }:
+            raise AttributeError
+
+        try:
+            cls_decl = self.__egg_decls__._classes[self.__egg_tp__.ident]
+        except Exception as e:
+            e.add_note(f"Error processing class {self.__egg_tp__.ident}")
+            raise
+
+        preserved_methods = cls_decl.preserved_methods
+        if name in preserved_methods:
+            return preserved_methods[name]
+
+        # if this is a class variable, return an expr for it, otherwise, assume it's a method
+        if name in cls_decl.class_variables:
+            return_tp = cls_decl.class_variables[name]
+            return RuntimeExpr(
+                self.__egg_decls_thunk__,
+                Thunk.value(TypedExprDecl(return_tp.type_ref, CallDecl(ClassVariableRef(self.__egg_tp__.ident, name)))),
+            )
+        if name in cls_decl.class_methods:
+            return RuntimeFunction(
+                self.__egg_decls_thunk__,
+                Thunk.value(ClassMethodRef(self.__egg_tp__.ident, name)),
+                self.__egg_tp__.to_just(),
+            )
+        # allow referencing properties and methods as class variables as well
+        if name in cls_decl.properties:
+            return RuntimeFunction(self.__egg_decls_thunk__, Thunk.value(PropertyRef(self.__egg_tp__.ident, name)))
+        if name in cls_decl.methods:
+            return RuntimeFunction(self.__egg_decls_thunk__, Thunk.value(MethodRef(self.__egg_tp__.ident, name)))
+
+        msg = f"Class {self.__egg_tp__.ident} has no method {name}"
+        raise AttributeError(msg) from None
+
+    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(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
+
+    @property
+    def __parameters__(self) -> tuple[object, ...]:
+        """
+        Emit a number of typevar params so that when using generic type aliases, we know how to resolve these properly.
+        """
+        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__.ident].match_args
+
+    @property
+    def __doc__(self) -> str | None:  # type: ignore[override]
+        return self.__egg_decls__._classes[self.__egg_tp__.ident].doc
+
+    @property
+    def __dict__(self) -> dict[str, Any]:  # type: ignore[override]
+        """
+        Return the dict so that this works with things like `DocTestFinder`.
+        """
+        return {attr: getattr(self, attr) for attr in dir(self)}
+
+
+class RuntimeFunctionMeta(type):
+    # Override getatribute on class so that it if we call RuntimeFunction.__module__ we get this module
+    # but if we call it on an instance we get the module of the instance so that it works with doctest.
+    def __getattribute__(cls, name: str) -> Any:
+        if name == "__module__":
+            return __name__
+        return super().__getattribute__(name)
+
+
+@dataclass
+class RuntimeFunction(DelayedDeclerations, metaclass=RuntimeFunctionMeta):
+    __egg_ref_thunk__: Callable[[], CallableRef]
+    # 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
+
+    __get__: None = None
+
+    @property
+    def __module__(self) -> str | None:  # type: ignore[override]
+        ref = self.__egg_ref__
+        if isinstance(ref, UnnamedFunctionRef):
+            return None
+        return ref.ident.module
+
+    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  # 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
+        decls = self.__egg_decls__.copy()
+        # Special case function application bc we dont support variadic generics yet generally
+        if signature == "fn-app":
+            fn, *rest_args = args
+            args = tuple(rest_args)
+            assert not kwargs
+            assert isinstance(fn, RuntimeExpr)
+            decls.update(fn)
+            function_value = fn.__egg_typed_expr__
+            fn_tp = function_value.tp
+            assert fn_tp.ident == Ident.builtin("UnstableFn")
+            fn_return_tp, *fn_arg_tps = fn_tp.args
+            signature = FunctionSignature(
+                tuple(tp.to_var() for tp in fn_arg_tps),
+                tuple(f"_{i}" for i in range(len(fn_arg_tps))),
+                (None,) * len(fn_arg_tps),
+                fn_return_tp.to_var(),
+            )
+        else:
+            function_value = None
+            assert isinstance(signature, FunctionSignature)
+
+        # Turn all keyword args into positional args
+        py_signature = to_py_signature(signature, self.__egg_decls__, _egg_partial_function)
+        try:
+            bound = py_signature.bind(*args, **kwargs)
+        except TypeError as 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
+        args = bound.args
+
+        tcs = TypeConstraintSolver(decls)
+        bound_tp = (
+            None
+            if self.__egg_bound__ is None
+            else self.__egg_bound__.__egg_typed_expr__.tp
+            if isinstance(self.__egg_bound__, RuntimeExpr)
+            else self.__egg_bound__
+        )
+        if (
+            bound_tp
+            and bound_tp.args
+            # Don't  bind class if we have a first class function arg, b/c we don't support that yet
+            and not function_value
+        ):
+            tcs.bind_class(bound_tp)
+        assert (operator.ge if signature.var_arg_type else operator.eq)(len(args), len(signature.arg_types))
+        cls_ident = bound_tp.ident if bound_tp else None
+        upcasted_args = [
+            resolve_literal(cast("TypeOrVarRef", tp), arg, Thunk.value(decls), tcs=tcs, cls_ident=cls_ident)
+            for arg, tp in zip_longest(args, signature.arg_types, fillvalue=signature.var_arg_type)
+        ]
+        decls.update(*upcasted_args)
+        arg_exprs = tuple(arg.__egg_typed_expr__ for arg in upcasted_args)
+        return_tp = tcs.substitute_typevars(signature.semantic_return_type, cls_ident)
+        bound_params = (
+            cast("JustTypeRef", bound_tp).args if isinstance(self.__egg_ref__, ClassMethodRef | InitRef) else ()
+        )
+        # If we were using unstable-app to call a funciton, add that function back as the first arg.
+        if function_value:
+            arg_exprs = (function_value, *arg_exprs)
+        expr_decl = CallDecl(self.__egg_ref__, arg_exprs, bound_params)
+        typed_expr_decl = TypedExprDecl(return_tp, expr_decl)
+        # If there is not return type, we are mutating the first arg
+        if not signature.return_type:
+            first_arg = upcasted_args[0]
+            first_arg.__egg_decls_thunk__ = Thunk.value(decls)
+            first_arg.__egg_typed_expr_thunk__ = Thunk.value(typed_expr_decl)
+            return None
+        return RuntimeExpr.__from_values__(decls, typed_expr_decl)
+
+    def __str__(self) -> str:
+        first_arg, bound_tp_params = None, None
+        match self.__egg_bound__:
+            case RuntimeExpr(_):
+                first_arg = self.__egg_bound__.__egg_typed_expr__.expr
+            case JustTypeRef(_, args):
+                bound_tp_params = args
+        return pretty_callable_ref(self.__egg_decls__, self.__egg_ref__, first_arg, bound_tp_params)
+
+    def __repr__(self) -> str:
+        return str(self)
+
+    @property
+    def __doc__(self) -> str | None:  # type: ignore[override]
+        decl = self.__egg_decls__.get_callable_decl(self.__egg_ref__)
+        if isinstance(decl, FunctionDecl | ConstructorDecl):
+            return decl.doc
+        return None
+
+
+def to_py_signature(sig: FunctionSignature, decls: Declarations, optional_args: bool) -> Signature:
+    """
+    Convert to a Python signature.
+
+    If optional_args is true, then all args will be treated as optional, as if a default was provided that makes them
+    a var with that arg name as the value.
+
+    Used for partial application to try binding a function with only some of its args.
+    """
+    parameters = [
+        Parameter(
+            n,
+            Parameter.POSITIONAL_OR_KEYWORD,
+            default=RuntimeExpr.__from_values__(decls, TypedExprDecl(t.to_just(), d or LetRefDecl(n)))
+            if d is not None or optional_args
+            else Parameter.empty,
+        )
+        for n, d, t in zip(sig.arg_names, sig.arg_defaults, sig.arg_types, strict=True)
+    ]
+    if isinstance(sig, FunctionSignature) and sig.var_arg_type is not None:
+        parameters.append(Parameter("__rest", Parameter.VAR_POSITIONAL))
+    return Signature(parameters)
+
+
+ON_CREATE_EXPR: None | Callable[[Callable[[], TypedExprDecl]], None] = None
+
+
+@dataclass
+class RuntimeExpr(DelayedDeclerations):
+    __egg_typed_expr_thunk__: Callable[[], TypedExprDecl]
+
+    def __post_init__(self) -> None:
+        if ON_CREATE_EXPR:
+            ON_CREATE_EXPR(self.__egg_typed_expr_thunk__)
+
+    @classmethod
+    def __from_values__(cls, d: Declarations, e: TypedExprDecl) -> RuntimeExpr:
+        return cls(Thunk.value(d), Thunk.value(e))
+
+    def __with_expr__(self, e: TypedExprDecl) -> RuntimeExpr:
+        return RuntimeExpr(self.__egg_decls_thunk__, Thunk.value(e))
+
+    @property
+    def __egg_typed_expr__(self) -> TypedExprDecl:
+        return self.__egg_typed_expr_thunk__()
+
+    def __getattr__(self, name: str) -> object:
+        if (method := _get_expr_method(self, name)) is not _no_method_sentinel:
+            return method
+        if name in self.__egg_class_decl__.properties:
+            fn = RuntimeFunction(
+                self.__egg_decls_thunk__, Thunk.value(PropertyRef(self.__egg_class_ident__, name)), self
+            )
+            return fn()
+        if (getattr_method := _get_expr_method(self, "__getattr__")) is not _no_method_sentinel:
+            return cast("Callable", getattr_method)(name)
+        raise AttributeError(f"{self.__egg_class_ident__} has no method {name}") from None
+
+    def __repr__(self) -> str:
+        """
+        The repr of the expr is the pretty printed version of the expr.
+        """
+        if (method := _get_expr_method(self, "__repr__")) is not _no_method_sentinel:
+            return cast("str", cast("Any", cast("Callable", method)()))
+        return str(self)
+
+    def __str__(self) -> str:
+        if (method := _get_expr_method(self, "__str__")) is not _no_method_sentinel:
+            return cast("str", cast("Any", cast("Callable", method)()))
+        return self.__egg_pretty__(None)
+
+    def __egg_pretty__(self, wrapping_fn: str | None) -> str:
+        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  # noqa: PLC0415
+
+        display(Code(str(self), language="python"))
+
+    def __dir__(self) -> Iterable[str]:
+        cls_decl = self.__egg_class_decl__
+        return (
+            list(cls_decl.class_methods)
+            + list(cls_decl.class_variables)
+            + list(cls_decl.preserved_methods)
+            + list(cls_decl.methods)
+            + list(cls_decl.properties)
+        )
+
+    @property
+    def __egg_class_ident__(self) -> Ident:
+        return self.__egg_typed_expr__.tp.ident
+
+    @property
+    def __egg_class_decl__(self) -> ClassDecl:
+        return self.__egg_decls__.get_class_decl(self.__egg_class_ident__)
+
+    # 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
+
+    def __getstate__(self) -> tuple[Declarations, TypedExprDecl]:
+        return self.__egg_decls__, self.__egg_typed_expr__
+
+    def __setstate__(self, d: tuple[Declarations, TypedExprDecl]) -> None:
+        self.__egg_decls_thunk__ = Thunk.value(d[0])
+        self.__egg_typed_expr_thunk__ = Thunk.value(d[1])
+
+    def __hash__(self) -> int:
+        if (method := _get_expr_method(self, "__hash__")) is not _no_method_sentinel:
+            return cast("int", cast("Any", cast("Callable", 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 _no_method_sentinel:
+            return cast("Callable", method)(other)
+
+        if not (isinstance(self, RuntimeExpr) and isinstance(other, RuntimeExpr)):
+            return NotImplemented
+        if self.__egg_typed_expr__.tp != other.__egg_typed_expr__.tp:
+            return NotImplemented
+
+        from .egraph import Fact  # noqa: PLC0415
+
+        return Fact(
+            Declarations.create(self, other),
+            EqDecl(self.__egg_typed_expr__.tp, self.__egg_typed_expr__.expr, other.__egg_typed_expr__.expr),
+        )
+
+    def __ne__(self, other: object) -> object:  # type: ignore[override]
+        if (method := _get_expr_method(self, "__ne__")) is not _no_method_sentinel:
+            return cast("Callable", method)(other)
+
+        from .egraph import BaseExpr, ne  # noqa: PLC0415
+
+        return ne(cast("BaseExpr", self)).to(cast("BaseExpr", other))
+
+    def __call__(
+        self, *args: object, **kwargs: object
+    ) -> object:  # define it here only for type checking, it will be overriden below
+        ...
+
+    def __replace_expr__(self, new_expr: RuntimeExpr) -> None:
+        self.__egg_decls_thunk__ = new_expr.__egg_decls_thunk__
+        self.__egg_typed_expr_thunk__ = new_expr.__egg_typed_expr_thunk__
+
+
+# Return a sentinel value to differentiate between no method and property that returns None
+_no_method_sentinel = object()
+
+
+def _get_expr_method(expr: RuntimeExpr, name: str) -> object:
+    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_ident__, name)), expr)
+    return _no_method_sentinel
+
+
+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 overrding with `__getattr__` does not suffice,
+    like for NumPy's `__array_ufunc__`.
+    """
+
+    def _defined_method(self: RuntimeExpr, *args, __name: str = name, **kwargs) -> object:
+        fn = cast("Callable", _get_expr_method(self, __name))
+        if fn is _no_method_sentinel:
+            if __name == "__hash__":
+                return hash(self.__egg_typed_expr__)
+            raise AttributeError(f"{self.__egg_class_ident__} expression has no method {__name}")
+        return fn(*args, **kwargs)
+
+    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)):
+    method_name = f"__r{name[2:]}" if r_method else name
+
+    def _numeric_binary_method(
+        self: object, other: object, name: str = name, r_method: bool = r_method, method_name: str = method_name
+    ) -> 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.
+        """
+        from .conversion import (  # noqa: PLC0415
+            ConvertError,
+            convert_to_same_type,
+            min_binary_conversion,
+            resolve_type,
+        )
+
+        # 1. switch if reversed method
+        if r_method:
+            self, other = other, self
+        # First check if we have a preserved method for this:
+        if isinstance(self, RuntimeExpr) and (
+            (preserved_method := self.__egg_class_decl__.preserved_methods.get(name)) is not None
+        ):
+            return preserved_method.__get__(self)(other)
+        # Then check if the self is a Python type and the other side has a preserved method
+        if (
+            not isinstance(self, RuntimeExpr)
+            and isinstance(other, RuntimeExpr)
+            and ((preserved_method := other.__egg_class_decl__.preserved_methods.get(name)) is not None)
+        ):
+            try:
+                new_self = convert_to_same_type(self, other)
+            except ConvertError:
+                pass
+            else:
+                return preserved_method.__get__(new_self)(other)
+
+        # 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
+                )
+            )
+        ):
+            best_method = min_binary_conversion(name, resolve_type(self), resolve_type(other))
+
+            if not best_method:
+                raise RuntimeError(
+                    f"Cannot resolve {name} for {resolve_type(self)} and {resolve_type(other)}, no conversion found"
+                )
+            self, other = best_method[0](self), best_method[1](other)
+
+        method_ref = MethodRef(self.__egg_class_ident__, name)
+        fn = RuntimeFunction(Thunk.value(self.__egg_decls__), Thunk.value(method_ref), self)
+        return fn(other)
+
+    setattr(RuntimeExpr, method_name, _numeric_binary_method)
+
+
+def resolve_callable(callable: object) -> tuple[CallableRef, Declarations]:
+    """
+    Resolves a runtime callable into a ref
+    """
+    # TODO: Make runtime class work with __match_args__
+    if isinstance(callable, RuntimeClass):
+        return InitRef(callable.__egg_tp__.ident), callable.__egg_decls__
+    match callable:
+        case RuntimeFunction(decls, ref, _):
+            return ref(), decls()
+        case RuntimeExpr(decl_thunk, expr_thunk):
+            if not isinstance((expr := expr_thunk().expr), CallDecl) or not isinstance(
+                expr.callable, ConstantRef | ClassVariableRef
+            ):
+                raise NotImplementedError(f"Can only turn constants or classvars into callable refs, not {expr}")
+            return expr.callable, decl_thunk()
+        case types.MethodWrapperType() if isinstance((slf := callable.__self__), RuntimeClass):
+            return MethodRef(slf.__egg_tp__.ident, callable.__name__), slf.__egg_decls__
+        case _:
+            raise NotImplementedError(f"Cannot turn {callable} of type {type(callable)} into a callable ref")
+
+
+def create_callable(decls: Declarations, ref: CallableRef) -> RuntimeClass | RuntimeFunction | RuntimeExpr:
+    """
+    Creates a callable object from a callable ref. This might not actually be callable, if the ref is a constant
+    or classvar then it is a value
+    """
+    match ref:
+        case InitRef(name):
+            return RuntimeClass(Thunk.value(decls), TypeRefWithVars(name))
+        case FunctionRef() | MethodRef() | ClassMethodRef() | PropertyRef() | UnnamedFunctionRef():
+            bound = JustTypeRef(ref.ident) if isinstance(ref, ClassMethodRef) else None
+            return RuntimeFunction(Thunk.value(decls), Thunk.value(ref), bound)
+        case ConstantRef(name):
+            tp = decls._constants[name].type_ref
+        case ClassVariableRef(cls_name, var_name):
+            tp = decls._classes[cls_name].class_variables[var_name].type_ref
+        case _:
+            assert_never(ref)
+    return RuntimeExpr.__from_values__(decls, TypedExprDecl(tp, CallDecl(ref)))