effectful 0.0.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

effectful/ops/semantics.py

@@ -1,24 +1,22 @@
+import collections.abc
 import contextlib
-import functools
-from typing import Any, Callable, Optional, Set, Type, TypeVar
+import dataclasses
+import types
+import typing
+from typing import Any
 
-import tree
-from typing_extensions import ParamSpec
-
-from effectful.ops.syntax import NoDefaultRule, deffn, defop, defterm
-from effectful.ops.types import Expr, Interpretation, Operation, Term
-
-P = ParamSpec("P")
-Q = ParamSpec("Q")
-S = TypeVar("S")
-T = TypeVar("T")
-V = TypeVar("V")
+from effectful.ops.syntax import defop
+from effectful.ops.types import (
+    Expr,
+    Interpretation,
+    NotHandled,  # noqa: F401
+    Operation,
+    Term,
+)
 
 
 @defop  # type: ignore
-def apply(
-    intp: Interpretation[S, T], op: Operation[P, S], *args: P.args, **kwargs: P.kwargs
-) -> T:
+def apply[**P, T](op: Operation[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
     """Apply ``op`` to ``args``, ``kwargs`` in interpretation ``intp``.
 
     Handling :func:`apply` changes the evaluation strategy of terms.
@@ -39,44 +37,25 @@ def apply(
 
     By installing an :func:`apply` handler, we capture the term instead:
 
-    >>> with handler({apply: lambda _, op, *args, **kwargs: op.__free_rule__(*args, **kwargs) }):
+    >>> def default(*args, **kwargs):
+    ...     raise NotHandled
+    >>> with handler({apply: default }):
     ...     term = mul(add(1, 2), 3)
-    >>> term
+    >>> print(str(term))
     mul(add(1, 2), 3)
 
     """
+    from effectful.internals.runtime import get_interpretation
+
+    intp = get_interpretation()
     if op in intp:
         return intp[op](*args, **kwargs)
     elif apply in intp:
-        return intp[apply](intp, op, *args, **kwargs)
+        return intp[apply](op, *args, **kwargs)
     else:
         return op.__default_rule__(*args, **kwargs)  # type: ignore
 
 
-@defop  # type: ignore
-def call(fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
-    """An operation that eliminates a callable term.
-
-    This operation is invoked by the ``__call__`` method of a callable term.
-
-    """
-    if not isinstance(fn, Term):
-        fn = defterm(fn)
-
-    if isinstance(fn, Term) and fn.op is deffn:
-        body: Expr[Callable[P, T]] = fn.args[0]
-        argvars: tuple[Operation, ...] = fn.args[1:]
-        kwvars: dict[str, Operation] = fn.kwargs
-        subs = {
-            **{v: functools.partial(lambda x: x, a) for v, a in zip(argvars, args)},
-            **{kwvars[k]: functools.partial(lambda x: x, kwargs[k]) for k in kwargs},
-        }
-        with handler(subs):
-            return evaluate(body)
-    else:
-        raise NoDefaultRule
-
-
 @defop
 def fwd(*args, **kwargs) -> Any:
     """Forward execution to the next most enclosing handler.
@@ -93,9 +72,7 @@ def fwd(*args, **kwargs) -> Any:
     raise RuntimeError("fwd should only be called in the context of a handler")
 
 
-def coproduct(
-    intp: Interpretation[S, T], intp2: Interpretation[S, T]
-) -> Interpretation[S, T]:
+def coproduct(intp: Interpretation, intp2: Interpretation) -> Interpretation:
     """The coproduct of two interpretations handles any effect that is handled
     by either. If both interpretations handle an effect, ``intp2`` takes
     precedence.
@@ -151,7 +128,7 @@ def coproduct(
         if op is fwd or op is _get_args:
             res[op] = i2  # fast path for special cases, should be equivalent if removed
         else:
-            i1 = intp.get(op, op.__default_rule__)  # type: ignore
+            i1 = intp.get(op, op.__default_rule__)
 
             # calling fwd in the right handler should dispatch to the left handler
             res[op] = _set_prompt(fwd, _restore_args(_save_args(i1)), _save_args(i2))
@@ -159,9 +136,7 @@ def coproduct(
     return res
 
 
-def product(
-    intp: Interpretation[S, T], intp2: Interpretation[S, T]
-) -> Interpretation[S, T]:
+def product(intp: Interpretation, intp2: Interpretation) -> Interpretation:
     """The product of two interpretations handles any effect that is handled by
     ``intp2``. Handlers in ``intp2`` may override handlers in ``intp``, but
     those changes are not visible to the handlers in ``intp``. In this way,
@@ -197,7 +172,7 @@ def product(
 
     """
     if any(op in intp for op in intp2):  # alpha-rename
-        renaming = {op: defop(op) for op in intp2 if op in intp}
+        renaming: Interpretation = {op: defop(op) for op in intp2 if op in intp}
         intp_fresh = {renaming.get(op, op): handler(renaming)(intp[op]) for op in intp}
         return product(intp_fresh, intp2)
     else:
@@ -207,7 +182,7 @@ def product(
 
 
 @contextlib.contextmanager
-def runner(intp: Interpretation[S, T]):
+def runner(intp: Interpretation):
     """Install an interpretation by taking a product with the current
     interpretation.
 
@@ -215,7 +190,7 @@ def runner(intp: Interpretation[S, T]):
     from effectful.internals.runtime import get_interpretation, interpreter
 
     @interpreter(get_interpretation())
-    def _reapply(_, op: Operation[P, S], *args: P.args, **kwargs: P.kwargs):
+    def _reapply[**P, S](op: Operation[P, S], *args: P.args, **kwargs: P.kwargs):
         return op(*args, **kwargs)
 
     with interpreter({apply: _reapply, **intp}):
@@ -223,7 +198,7 @@ def runner(intp: Interpretation[S, T]):
 
 
 @contextlib.contextmanager
-def handler(intp: Interpretation[S, T]):
+def handler(intp: Interpretation):
     """Install an interpretation by taking a coproduct with the current
     interpretation.
 
@@ -234,7 +209,7 @@ def handler(intp: Interpretation[S, T]):
         yield intp
 
 
-def evaluate(expr: Expr[T], *, intp: Optional[Interpretation[S, T]] = None) -> Expr[T]:
+def evaluate[T](expr: Expr[T], *, intp: Interpretation | None = None) -> Expr[T]:
     """Evaluate expression ``expr`` using interpretation ``intp``. If no
     interpretation is provided, uses the current interpretation.
 
@@ -245,33 +220,61 @@ def evaluate(expr: Expr[T], *, intp: Optional[Interpretation[S, T]] = None) -> E
 
     >>> @defop
     ... def add(x: int, y: int) -> int:
-    ...     raise NoDefaultRule
+    ...     raise NotHandled
     >>> expr = add(1, add(2, 3))
-    >>> expr
+    >>> print(str(expr))
     add(1, add(2, 3))
     >>> evaluate(expr, intp={add: lambda x, y: x + y})
     6
 
     """
-    if intp is None:
-        from effectful.internals.runtime import get_interpretation
-
-        intp = get_interpretation()
+    from effectful.internals.runtime import get_interpretation, interpreter
 
-    expr = defterm(expr) if not isinstance(expr, Term) else expr
+    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)
+        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, Operation):
+        op_intp = get_interpretation().get(expr, expr)
+        return op_intp if isinstance(op_intp, Operation) else expr  # type: ignore
+    elif isinstance(expr, collections.abc.Mapping):
+        if isinstance(expr, collections.defaultdict):
+            return type(expr)(expr.default_factory, evaluate(tuple(expr.items())))  # type: ignore
+        elif isinstance(expr, types.MappingProxyType):
+            return type(expr)(dict(evaluate(tuple(expr.items()))))  # type: ignore
+        else:
+            return type(expr)(evaluate(tuple(expr.items())))  # type: ignore
+    elif isinstance(expr, collections.abc.Sequence):
+        if isinstance(expr, str | bytes):
+            return typing.cast(T, expr)  # mypy doesnt like ignore here, so we use cast
+        else:
+            return type(expr)(evaluate(item) for item in expr)  # type: ignore
+    elif isinstance(expr, collections.abc.Set):
+        if isinstance(expr, collections.abc.ItemsView | collections.abc.KeysView):
+            return {evaluate(item) for item in expr}  # type: ignore
+        else:
+            return type(expr)(evaluate(item) for item in expr)  # type: ignore
+    elif isinstance(expr, collections.abc.ValuesView):
+        return [evaluate(item) for item in expr]  # type: ignore
+    elif dataclasses.is_dataclass(expr) and not isinstance(expr, type):
+        return typing.cast(
+            T,
+            dataclasses.replace(
+                expr,
+                **{
+                    field.name: evaluate(getattr(expr, field.name))
+                    for field in dataclasses.fields(expr)
+                },
+            ),
         )
-        return apply.__default_rule__(intp, expr.op, *args, **kwargs)  # type: ignore
-    elif tree.is_nested(expr):
-        return tree.map_structure(functools.partial(evaluate, intp=intp), expr)
     else:
-        return expr
+        return typing.cast(T, expr)
 
 
-def typeof(term: Expr[T]) -> Type[T]:
+def typeof[T](term: Expr[T]) -> type[T]:
     """Return the type of an expression.
 
     **Example usage**:
@@ -280,46 +283,63 @@ def typeof(term: Expr[T]) -> Type[T]:
 
     >>> @defop
     ... def cmp(x: int, y: int) -> bool:
-    ...     raise NoDefaultRule
+    ...     raise NotHandled
     >>> typeof(cmp(1, 2))
     <class 'bool'>
 
     Types can be computed in the presence of type variables.
 
-    >>> from typing import TypeVar
-    >>> T = TypeVar('T')
     >>> @defop
-    ... def if_then_else(x: bool, a: T, b: T) -> T:
-    ...     raise NoDefaultRule
+    ... def if_then_else[T](x: bool, a: T, b: T) -> T:
+    ...     raise NotHandled
     >>> typeof(if_then_else(True, 0, 1))
     <class 'int'>
 
     """
     from effectful.internals.runtime import interpreter
 
-    with interpreter({apply: lambda _, op, *a, **k: op.__type_rule__(*a, **k)}):
-        return evaluate(term)  # type: ignore
+    with interpreter({apply: lambda op, *a, **k: op.__type_rule__(*a, **k)}):
+        if isinstance(term, Term):
+            # If term is a Term, we evaluate it to get its type
+            tp = evaluate(term)
+            if isinstance(tp, typing.TypeVar):
+                tp = (
+                    tp.__bound__
+                    if tp.__bound__
+                    else tp.__constraints__[0]
+                    if tp.__constraints__
+                    else object
+                )
+            if isinstance(tp, types.UnionType):
+                raise TypeError(
+                    f"Cannot determine type of {term} because it is a union type: {tp}"
+                )
+            return typing.get_origin(tp) or tp  # type: ignore
+        else:
+            return type(term)
 
 
-def fvsof(term: Expr[S]) -> Set[Operation]:
+def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]:
     """Return the free variables of an expression.
 
     **Example usage**:
 
     >>> @defop
     ... def f(x: int, y: int) -> int:
-    ...     raise NoDefaultRule
-    >>> fvsof(f(1, 2))
-    {f}
-
+    ...     raise NotHandled
+    >>> fvs = fvsof(f(1, 2))
+    >>> assert f in fvs
+    >>> assert len(fvs) == 1
     """
     from effectful.internals.runtime import interpreter
 
-    _fvs: Set[Operation] = set()
+    _fvs: set[Operation] = set()
 
-    def _update_fvs(_, op, *args, **kwargs):
+    def _update_fvs(op, *args, **kwargs):
         _fvs.add(op)
-        for bound_var in op.__fvs_rule__(*args, **kwargs):
+        bindings = op.__fvs_rule__(*args, **kwargs)
+        for bound_var in set().union(*(*bindings.args, *bindings.kwargs.values())):
+            assert isinstance(bound_var, Operation)
             if bound_var in _fvs:
                 _fvs.remove(bound_var)