Add defdata patterns and Operations for polymorphic collection and Number types

effectful/handlers/numbers.py

@@ -3,244 +3,28 @@
 """
 
 import numbers
-import operator
-from typing import Any
 
-from effectful.ops.syntax import defdata, defop
-from effectful.ops.types import Expr, Operation, Term
-
-
-@defdata.register(numbers.Number)
-@numbers.Number.register
-class _NumberTerm(Term[numbers.Number]):
-    def __init__(
-        self, op: Operation[..., numbers.Number], *args: Expr, **kwargs: Expr
-    ) -> None:
-        self._op = op
-        self._args = args
-        self._kwargs = kwargs
-
-    @property
-    def op(self) -> Operation[..., numbers.Number]:
-        return self._op
-
-    @property
-    def args(self) -> tuple:
-        return self._args
-
-    @property
-    def kwargs(self) -> dict:
-        return self._kwargs
-
-    def __hash__(self):
-        return hash((self.op, tuple(self.args), tuple(self.kwargs.items())))
-
-
-# Complex specific methods
-def _wrap_cmp(op):
-    def _wrapped_op[T_Number: numbers.Number](x: T_Number, y: T_Number) -> bool:
-        if not any(isinstance(a, Term) for a in (x, y)):
-            return op(x, y)
-        else:
-            raise NotImplementedError
-
-    _wrapped_op.__name__ = op.__name__
-    return _wrapped_op
-
-
-def _wrap_binop(op):
-    def _wrapped_op[T_Number: numbers.Number](x: T_Number, y: T_Number) -> T_Number:
-        if not any(isinstance(a, Term) for a in (x, y)):
-            return op(x, y)
-        else:
-            raise NotImplementedError
-
-    _wrapped_op.__name__ = op.__name__
-    return _wrapped_op
-
-
-def _wrap_unop(op):
-    def _wrapped_op[T_Number: numbers.Number](x: T_Number) -> T_Number:
-        if not isinstance(x, Term):
-            return op(x)
-        else:
-            raise NotImplementedError
-
-    _wrapped_op.__name__ = op.__name__
-    return _wrapped_op
-
-
-add = defop(_wrap_binop(operator.add))
-neg = defop(_wrap_unop(operator.neg))
-pos = defop(_wrap_unop(operator.pos))
-sub = defop(_wrap_binop(operator.sub))
-mul = defop(_wrap_binop(operator.mul))
-truediv = defop(_wrap_binop(operator.truediv))
-pow = defop(_wrap_binop(operator.pow))
-abs = defop(_wrap_unop(operator.abs))
-eq = defop(_wrap_cmp(operator.eq))
-
-
-@defdata.register(numbers.Complex)
-@numbers.Complex.register
-class _ComplexTerm(_NumberTerm, Term[numbers.Complex]):
-    def __bool__(self) -> bool:
-        raise ValueError("Cannot convert term to bool")
-
-    def __add__(self, other: Any) -> numbers.Real:
-        return add(self, other)
-
-    def __radd__(self, other: Any) -> numbers.Real:
-        return add(other, self)
-
-    def __neg__(self):
-        return neg(self)
-
-    def __pos__(self):
-        return pos(self)
-
-    def __sub__(self, other: Any) -> numbers.Real:
-        return sub(self, other)
-
-    def __rsub__(self, other: Any) -> numbers.Real:
-        return sub(other, self)
-
-    def __mul__(self, other: Any) -> numbers.Real:
-        return mul(self, other)
-
-    def __rmul__(self, other: Any) -> numbers.Real:
-        return mul(other, self)
-
-    def __truediv__(self, other: Any) -> numbers.Real:
-        return truediv(self, other)
-
-    def __rtruediv__(self, other: Any) -> numbers.Real:
-        return truediv(other, self)
-
-    def __pow__(self, other: Any) -> numbers.Real:
-        return pow(self, other)
-
-    def __rpow__(self, other: Any) -> numbers.Real:
-        return pow(other, self)
-
-    def __abs__(self) -> numbers.Real:
-        return abs(self)
-
-    def __eq__(self, other: Any) -> bool:
-        return eq(self, other)
-
-
-# Real specific methods
-floordiv = defop(_wrap_binop(operator.floordiv))
-mod = defop(_wrap_binop(operator.mod))
-lt = defop(_wrap_cmp(operator.lt))
-le = defop(_wrap_cmp(operator.le))
-gt = defop(_wrap_cmp(operator.gt))
-ge = defop(_wrap_cmp(operator.ge))
-
-
-@defdata.register(numbers.Real)
-@numbers.Real.register
-class _RealTerm(_ComplexTerm, Term[numbers.Real]):
-    # Real specific methods
-    def __float__(self) -> float:
-        raise ValueError("Cannot convert term to float")
-
-    def __trunc__(self) -> numbers.Integral:
-        raise NotImplementedError
-
-    def __floor__(self) -> numbers.Integral:
-        raise NotImplementedError
-
-    def __ceil__(self) -> numbers.Integral:
-        raise NotImplementedError
-
-    def __round__(self, ndigits=None) -> numbers.Integral:
-        raise NotImplementedError
-
-    def __floordiv__(self, other):
-        return floordiv(self, other)
-
-    def __rfloordiv__(self, other):
-        return floordiv(other, self)
-
-    def __mod__(self, other):
-        return mod(self, other)
-
-    def __rmod__(self, other):
-        return mod(other, self)
-
-    def __lt__(self, other):
-        return lt(self, other)
-
-    def __le__(self, other):
-        return le(self, other)
-
-
-@defdata.register(numbers.Rational)
-@numbers.Rational.register
-class _RationalTerm(_RealTerm, Term[numbers.Rational]):
-    @property
-    def numerator(self):
-        raise NotImplementedError
-
-    @property
-    def denominator(self):
-        raise NotImplementedError
-
-
-# Integral specific methods
-index = defop(_wrap_unop(operator.index))
-lshift = defop(_wrap_binop(operator.lshift))
-rshift = defop(_wrap_binop(operator.rshift))
-and_ = defop(_wrap_binop(operator.and_))
-xor = defop(_wrap_binop(operator.xor))
-or_ = defop(_wrap_binop(operator.or_))
-invert = defop(_wrap_unop(operator.invert))
-
-
-@defdata.register(numbers.Integral)
-@numbers.Integral.register
-class _IntegralTerm(_RationalTerm, Term[numbers.Integral]):
-    # Integral specific methods
-    def __int__(self) -> int:
-        raise ValueError("Cannot convert term to int")
-
-    def __index__(self) -> numbers.Integral:
-        return index(self)
-
-    def __pow__(self, exponent: Any, modulus=None) -> numbers.Integral:
-        return pow(self, exponent)
-
-    def __lshift__(self, other):
-        return lshift(self, other)
-
-    def __rlshift__(self, other):
-        return lshift(other, self)
-
-    def __rshift__(self, other):
-        return rshift(self, other)
-
-    def __rrshift__(self, other):
-        return rshift(other, self)
-
-    def __and__(self, other):
-        return and_(self, other)
-
-    def __rand__(self, other):
-        return and_(other, self)
-
-    def __xor__(self, other):
-        return xor(self, other)
-
-    def __rxor__(self, other):
-        return xor(other, self)
-
-    def __or__(self, other):
-        return or_(self, other)
-
-    def __ror__(self, other):
-        return or_(other, self)
-
-    def __invert__(self):
-        return invert(self)
+from effectful.ops.syntax import defdata
+
+add = defdata.dispatch(numbers.Number).__add__
+neg = defdata.dispatch(numbers.Number).__neg__
+pos = defdata.dispatch(numbers.Number).__pos__
+sub = defdata.dispatch(numbers.Number).__sub__
+mul = defdata.dispatch(numbers.Number).__mul__
+truediv = defdata.dispatch(numbers.Number).__truediv__
+pow = defdata.dispatch(numbers.Number).__pow__
+abs = defdata.dispatch(numbers.Number).__abs__
+floordiv = defdata.dispatch(numbers.Number).__floordiv__
+mod = defdata.dispatch(numbers.Number).__mod__
+eq = defdata.dispatch(numbers.Number).__eq__
+lt = defdata.dispatch(numbers.Number).__lt__
+le = defdata.dispatch(numbers.Number).__le__
+gt = defdata.dispatch(numbers.Number).__gt__
+ge = defdata.dispatch(numbers.Number).__ge__
+index = defdata.dispatch(numbers.Number).__index__
+lshift = defdata.dispatch(numbers.Number).__lshift__
+rshift = defdata.dispatch(numbers.Number).__rshift__
+and_ = defdata.dispatch(numbers.Number).__and__
+xor = defdata.dispatch(numbers.Number).__xor__
+or_ = defdata.dispatch(numbers.Number).__or__
+invert = defdata.dispatch(numbers.Number).__invert__

effectful/ops/semantics.py

@@ -1,12 +1,11 @@
+import collections.abc
 import contextlib
 import functools
 import types
 import typing
 from collections.abc import Callable
 from typing import Any
 
-import tree
-
 from effectful.ops.syntax import deffn, defop
 from effectful.ops.types import Expr, Interpretation, Operation, Term
 
@@ -68,7 +67,7 @@ def call[**P, T](fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
         }
         with handler(subs):
             return evaluate(body)
-    elif not any(isinstance(a, Term) for a in tree.flatten((fn, args, kwargs))):
+    elif not any(isinstance(a, Term) for a in (fn, *args, *kwargs.values())):
         return fn(*args, **kwargs)
     else:
         raise NotImplementedError
@@ -246,18 +245,24 @@ def evaluate[T](expr: Expr[T], *, intp: Interpretation | None = None) -> Expr[T]
     6
 
     """
-    if intp is None:
-        from effectful.internals.runtime import get_interpretation
+    from effectful.internals.runtime import get_interpretation, interpreter
 
-        intp = get_interpretation()
+    if intp is not None:
+        return interpreter(intp)(evaluate)(expr)
 
     if isinstance(expr, Term):
-        (args, kwargs) = tree.map_structure(
-            functools.partial(evaluate, intp=intp), (expr.args, expr.kwargs)
-        )
-        return apply.__default_rule__(intp, expr.op, *args, **kwargs)
-    elif tree.is_nested(expr):
-        return tree.map_structure(functools.partial(evaluate, intp=intp), expr)
+        args = tuple(evaluate(arg) for arg in expr.args)
+        kwargs = {k: evaluate(v) for k, v in expr.kwargs.items()}
+        return expr.op(*args, **kwargs)
+    elif isinstance(expr, collections.abc.Mapping):
+        return type(expr)((evaluate(k), evaluate(v)) for k, v in expr.items())
+    elif isinstance(expr, collections.abc.Sequence) and not isinstance(expr, str):
+        return type(expr)(evaluate(e) for e in expr)
+    elif isinstance(expr, collections.abc.Set):
+        return type(expr)(evaluate(e) for e in expr)
+    elif isinstance(expr, Operation):
+        op_intp = get_interpretation().get(expr, expr)
+        return op_intp if isinstance(op_intp, Operation) else expr
     else:
         return expr