egglog-python/python/egglog/runtime.py at main · egraphs-good/egglog-python

914 lines (766 loc) · 35.8 KB
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, field, replace
from inspect import Parameter, Signature
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",
    "DUMMY_VALUE",
    "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 situations where that is more ergonomic.
    new_decls, tp = resolve_type_annotation(tp)
    decls |= new_decls
    return tp
def resolve_type_annotation(tp: object) -> tuple[DeclarationsLike, 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, TypeVarRef.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, TypeVarRef):
        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.
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")
        return RuntimeClass.__getattr__(owner, self.name)
RUNTIME_CLASS_DESCRIPTORS: dict[str, RuntimeClassDescriptor] = {
    method: RuntimeClassDescriptor(method) for method in CLASS_DESCRIPTOR_METHODS
@dataclass(match_args=False)
class RuntimeClass(DelayedDeclarations, 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
    __egg_attr_cache__: dict[str, RuntimeFunction | RuntimeExpr | Callable] = field(
        init=False, repr=False, default_factory=dict
    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_function_types=self.__egg_tp__.args)
            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"
            # Clip off the remaining arguments
            captured_args = len(partial_args) + int(
                isinstance(fn_arg, RuntimeFunction) and isinstance(fn_arg.__egg_bound__, RuntimeExpr)
            value = PartialCallDecl(replace(call, args=call.args[:captured_args]))
            remaining_arg_types = [a.tp for a in call.args[captured_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[DeclarationsLike, ...], 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 should replace all existing args of typevars with new
        # args
        if old_args := self.__egg_tp__.args:
            is_typevar = [isinstance(arg, TypeVarRef) 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:  # noqa: C901
        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:
            return self.__egg_attr_cache__[name]
        except KeyError:
            pass
        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:
            res = preserved_methods[name]
        # if this is a class variable, return an expr for it, otherwise, assume it's a method
        elif name in cls_decl.class_variables:
            return_tp = cls_decl.class_variables[name]
            res = RuntimeExpr(
                self.__egg_decls_thunk__,
                Thunk.value(TypedExprDecl(return_tp.type_ref, CallDecl(ClassVariableRef(self.__egg_tp__.ident, name)))),
        else:
            if name in cls_decl.class_methods:
                callable_ref: CallableRef = ClassMethodRef(self.__egg_tp__.ident, name)
            # allow referencing properties and methods as class variables as well
            elif name in cls_decl.properties:
                callable_ref = PropertyRef(self.__egg_tp__.ident, name)
            elif name in cls_decl.methods:
                callable_ref = MethodRef(self.__egg_tp__.ident, name)
            else:
                msg = f"Class {self.__egg_tp__.ident} has no method {name}"
                raise AttributeError(msg) from None
            res = RuntimeFunction(
                self.__egg_decls_thunk__,
                Thunk.value(callable_ref),
                self.__egg_tp__.to_just() if self.__egg_tp__.args else None,
        self.__egg_attr_cache__[name] = res
        return res
    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)
class RuntimeFunction(DelayedDeclarations, metaclass=RuntimeFunctionMeta):
    __egg_ref_thunk__: Callable[[], CallableRef]
    # Either they bound class for something like `Vec[Int].create` or a RuntimeExpr for bound methods
    # 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__(  # noqa: C901,PLR0912
        self, *args: object, _egg_function_types: tuple[TypeOrVarRef, ...] | None = None, **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__, optional_args=_egg_function_types is not None)
        try:
            bound = py_signature.bind(*args, **kwargs)
        except TypeError as err:
            err.add_note(f"when calling {self} with args {args} and kwargs {kwargs}")
            raise
        del kwargs
        bound.apply_defaults()
        assert not bound.kwargs
        args = bound.args
        tcs = TypeConstraintSolver()
        if isinstance(self.__egg_bound__, JustTypeRef) and self.__egg_bound__.args:
            if function_value:
                msg = "Cannot have both bound type params and function value"
                raise ValueError(msg)
            tcs.bind_class(self.__egg_bound__, decls)
            bound_tp_params = self.__egg_bound__.args
        else:
            bound_tp_params = ()
        assert (operator.ge if signature.var_arg_type else operator.eq)(len(args), len(signature.arg_types))
        # Hack to allow being explicit on function types when casting. # noqa: FIX004
        for _fn_tp in _egg_function_types or ():
            try:
                _fn_tp_just = _fn_tp.to_just()
            except TypeVarError:
                continue
            tcs.bind_class(_fn_tp_just, decls)
            if _fn_tp_just.args:
        # Try using any runtime expressions passed in to help infer typevars
        for arg, tp in zip(args, signature.all_args, strict=False):
            if not isinstance(arg, RuntimeExpr):
                continue
            try:
                tcs.infer_typevars(tp, arg.__egg_typed_expr__.tp)
            # If this leads to an incompatibility, just skip it, since it could need to be upcasted
            except TypeConstraintError:
                continue
        # Now at this point we should be able to resolve all the typevars
        upcasted_args = [
            resolve_literal(
                tcs.substitute_typevars_try_function(tp, arg, Thunk.value(decls)).to_var(), arg, Thunk.value(decls)
            for arg, tp in zip(args, signature.all_args, strict=False)
        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)
        # If we were using unstable-app to call a function, 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_tp_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
# sentinel that will be upcasted to any type with a DummyDecl value
DUMMY_VALUE = object()
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
    `DUMMY_VALUE`
    Used for partial application to try binding a function with only some of its args.
    parameters = [
        Parameter(
            Parameter.POSITIONAL_OR_KEYWORD,
            default=RuntimeExpr.__from_values__(decls, TypedExprDecl(t.to_just(), d))
            if d is not None
            else (DUMMY_VALUE if 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
class RuntimeExpr(DelayedDeclarations):
    __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 overriding 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:
            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():
            return RuntimeFunction(Thunk.value(decls), Thunk.value(ref), None)
        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)))
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

runtime.py

Latest commit

History

runtime.py

File metadata and controls
