effectful 0.0.1__py3-none-any.whl

effectful/internals/base_impl.py

@@ -0,0 +1,259 @@
+import collections
+import functools
+import inspect
+import typing
+from typing import Callable, Generic, Mapping, Sequence, Set, Type, TypeVar
+
+import tree
+from typing_extensions import ParamSpec
+
+from effectful.ops.types import Expr, Operation, Term
+
+P = ParamSpec("P")
+Q = ParamSpec("Q")
+S = TypeVar("S")
+T = TypeVar("T")
+V = TypeVar("V")
+
+
+def rename(
+    subs: Mapping[Operation[..., S], Operation[..., S]],
+    leaf_value: V,  # Union[Term[V], Operation[..., V], V],
+) -> V:  # Union[Term[V], Operation[..., V], V]:
+    from effectful.internals.runtime import interpreter
+    from effectful.ops.semantics import apply, evaluate
+
+    if isinstance(leaf_value, Operation):
+        return subs.get(leaf_value, leaf_value)  # type: ignore
+    elif isinstance(leaf_value, Term):
+        with interpreter(
+            {apply: lambda _, op, *a, **k: op.__free_rule__(*a, **k), **subs}
+        ):
+            return evaluate(leaf_value)  # type: ignore
+    else:
+        return leaf_value
+
+
+class _BaseOperation(Generic[Q, V], Operation[Q, V]):
+    signature: Callable[Q, V]
+
+    def __init__(self, signature: Callable[Q, V]):
+        functools.update_wrapper(self, signature)
+        self.signature = signature
+
+    def __eq__(self, other):
+        if not isinstance(other, Operation):
+            return NotImplemented
+        return self.signature == other.signature
+
+    def __hash__(self):
+        return hash(self.signature)
+
+    def __default_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> "Expr[V]":
+        from effectful.ops.syntax import NoDefaultRule
+
+        try:
+            return self.signature(*args, **kwargs)
+        except NoDefaultRule:
+            return self.__free_rule__(*args, **kwargs)
+
+    def __free_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> "Expr[V]":
+        from effectful.ops.syntax import Bound, Scoped, defdata, defop
+
+        sig = inspect.signature(self.signature)
+        bound_sig = sig.bind(*args, **kwargs)
+        bound_sig.apply_defaults()
+
+        bound_vars: dict[int, set[Operation]] = collections.defaultdict(set)
+        scoped_args: dict[int, set[str]] = collections.defaultdict(set)
+        unscoped_args: set[str] = set()
+        for param_name, param in bound_sig.signature.parameters.items():
+            if typing.get_origin(param.annotation) is typing.Annotated:
+                for anno in param.annotation.__metadata__:
+                    if isinstance(anno, Bound):
+                        scoped_args[anno.scope].add(param_name)
+                        if param.kind is inspect.Parameter.VAR_POSITIONAL:
+                            assert isinstance(bound_sig.arguments[param_name], tuple)
+                            for bound_var in bound_sig.arguments[param_name]:
+                                bound_vars[anno.scope].add(bound_var)
+                        elif param.kind is inspect.Parameter.VAR_KEYWORD:
+                            assert isinstance(bound_sig.arguments[param_name], dict)
+                            for bound_var in bound_sig.arguments[param_name].values():
+                                bound_vars[anno.scope].add(bound_var)
+                        else:
+                            bound_vars[anno.scope].add(bound_sig.arguments[param_name])
+                    elif isinstance(anno, Scoped):
+                        scoped_args[anno.scope].add(param_name)
+            else:
+                unscoped_args.add(param_name)
+
+        # TODO replace this temporary check with more general scope level propagation
+        if bound_vars:
+            min_scope = min(bound_vars.keys(), default=0)
+            scoped_args[min_scope] |= unscoped_args
+            max_scope = max(bound_vars.keys(), default=0)
+            assert all(s in bound_vars or s > max_scope for s in scoped_args.keys())
+
+        # recursively rename bound variables from innermost to outermost scope
+        for scope in sorted(bound_vars.keys()):
+            # create fresh variables for each bound variable in the scope
+            renaming_map = {var: defop(var) for var in bound_vars[scope]}
+            # get just the arguments that are in the scope
+            for name in scoped_args[scope]:
+                bound_sig.arguments[name] = tree.map_structure(
+                    lambda a: rename(renaming_map, a),
+                    bound_sig.arguments[name],
+                )
+
+        return defdata(_BaseTerm(self, bound_sig.args, bound_sig.kwargs))
+
+    def __type_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> Type[V]:
+        sig = inspect.signature(self.signature)
+        bound_sig = sig.bind(*args, **kwargs)
+        bound_sig.apply_defaults()
+
+        anno = sig.return_annotation
+        if anno is inspect.Signature.empty:
+            return typing.cast(Type[V], object)
+        elif isinstance(anno, typing.TypeVar):
+            # rudimentary but sound special-case type inference sufficient for syntax ops:
+            # if the return type annotation is a TypeVar,
+            # look for a parameter with the same annotation and return its type,
+            # otherwise give up and return Any/object
+            for name, param in bound_sig.signature.parameters.items():
+                if param.annotation is anno and param.kind not in (
+                    inspect.Parameter.VAR_POSITIONAL,
+                    inspect.Parameter.VAR_KEYWORD,
+                ):
+                    arg = bound_sig.arguments[name]
+                    tp: Type[V] = type(arg) if not isinstance(arg, type) else arg
+                    return tp
+            return typing.cast(Type[V], object)
+        elif typing.get_origin(anno) is typing.Annotated:
+            tp = typing.get_args(anno)[0]
+            if not typing.TYPE_CHECKING:
+                tp = tp if typing.get_origin(tp) is None else typing.get_origin(tp)
+            return tp
+        elif typing.get_origin(anno) is not None:
+            return typing.get_origin(anno)
+        else:
+            return anno
+
+    def __fvs_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> Set[Operation]:
+        from effectful.ops.syntax import Bound
+
+        sig = inspect.signature(self.signature)
+        bound_sig = sig.bind(*args, **kwargs)
+        bound_sig.apply_defaults()
+
+        bound_vars: Set[Operation] = set()
+        for param_name, param in bound_sig.signature.parameters.items():
+            if typing.get_origin(param.annotation) is typing.Annotated:
+                for anno in param.annotation.__metadata__:
+                    if isinstance(anno, Bound):
+                        if param.kind is inspect.Parameter.VAR_POSITIONAL:
+                            for bound_var in bound_sig.arguments[param_name]:
+                                bound_vars.add(bound_var)
+                        elif param.kind is inspect.Parameter.VAR_KEYWORD:
+                            for bound_var in bound_sig.arguments[param_name].values():
+                                bound_vars.add(bound_var)
+                        else:
+                            bound_var = bound_sig.arguments[param_name]
+                            bound_vars.add(bound_var)
+
+        return bound_vars
+
+    def __repr_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> str:
+        args_str = ", ".join(map(str, args)) if args else ""
+        kwargs_str = (
+            ", ".join(f"{k}={str(v)}" for k, v in kwargs.items()) if kwargs else ""
+        )
+
+        ret = f"{self.signature.__name__}({args_str}"
+        if kwargs:
+            ret += f"{', ' if args else ''}"
+        ret += f"{kwargs_str})"
+        return ret
+
+    def __repr__(self):
+        return self.signature.__name__
+
+
+class _BaseTerm(Generic[T], Term[T]):
+    _op: Operation[..., T]
+    _args: Sequence[Expr]
+    _kwargs: Mapping[str, Expr]
+
+    def __init__(
+        self,
+        op: Operation[..., T],
+        args: Sequence[Expr],
+        kwargs: Mapping[str, Expr],
+    ):
+        self._op = op
+        self._args = args
+        self._kwargs = kwargs
+
+    def __eq__(self, other) -> bool:
+        from effectful.ops.syntax import syntactic_eq
+
+        return syntactic_eq(self, other)
+
+    @property
+    def op(self):
+        return self._op
+
+    @property
+    def args(self):
+        return self._args
+
+    @property
+    def kwargs(self):
+        return self._kwargs
+
+
+class _CallableTerm(Generic[P, T], _BaseTerm[collections.abc.Callable[P, T]]):
+    def __call__(self, *args: Expr, **kwargs: Expr) -> Expr[T]:
+        from effectful.ops.semantics import call
+
+        return call(self, *args, **kwargs)  # type: ignore
+
+
+def _unembed_callable(value: Callable[P, T]) -> Expr[Callable[P, T]]:
+    from effectful.internals.runtime import interpreter
+    from effectful.ops.semantics import apply, call
+    from effectful.ops.syntax import deffn, defop
+
+    assert not isinstance(value, Term)
+
+    try:
+        sig = inspect.signature(value)
+    except ValueError:
+        return value
+
+    for name, param in sig.parameters.items():
+        if param.kind in (
+            inspect.Parameter.VAR_POSITIONAL,
+            inspect.Parameter.VAR_KEYWORD,
+        ):
+            raise NotImplementedError(
+                f"cannot unembed {value}: parameter {name} is variadic"
+            )
+
+    bound_sig = sig.bind(
+        **{name: defop(param.annotation) for name, param in sig.parameters.items()}
+    )
+    bound_sig.apply_defaults()
+
+    with interpreter(
+        {
+            apply: lambda _, op, *a, **k: op.__free_rule__(*a, **k),
+            call: call.__default_rule__,
+        }
+    ):
+        body = value(
+            *[a() for a in bound_sig.args],
+            **{k: v() for k, v in bound_sig.kwargs.items()},
+        )
+
+    return deffn(body, *bound_sig.args, **bound_sig.kwargs)

effectful/internals/runtime.py

@@ -0,0 +1,78 @@
+import contextlib
+import dataclasses
+import functools
+from typing import Callable, Generic, Mapping, Tuple, TypeVar
+
+from typing_extensions import ParamSpec
+
+from effectful.ops.syntax import defop
+from effectful.ops.types import Interpretation, Operation
+
+P = ParamSpec("P")
+S = TypeVar("S")
+T = TypeVar("T")
+
+
+@dataclasses.dataclass
+class Runtime(Generic[S, T]):
+    interpretation: "Interpretation[S, T]"
+
+
+@functools.lru_cache(maxsize=1)
+def get_runtime() -> Runtime:
+    return Runtime(interpretation={})
+
+
+def get_interpretation():
+    return get_runtime().interpretation
+
+
+@contextlib.contextmanager
+def interpreter(intp: "Interpretation"):
+
+    r = get_runtime()
+    old_intp = r.interpretation
+    try:
+        old_intp, r.interpretation = r.interpretation, dict(intp)
+        yield intp
+    finally:
+        r.interpretation = old_intp
+
+
+@defop
+def _get_args() -> Tuple[Tuple, Mapping]:
+    return ((), {})
+
+
+def _restore_args(fn: Callable[P, T]) -> Callable[P, T]:
+    @functools.wraps(fn)
+    def _cont_wrapper(*a: P.args, **k: P.kwargs) -> T:
+        a, k = (a, k) if a or k else _get_args()  # type: ignore
+        return fn(*a, **k)
+
+    return _cont_wrapper
+
+
+def _save_args(fn: Callable[P, T]) -> Callable[P, T]:
+    from effectful.ops.semantics import handler
+
+    @functools.wraps(fn)
+    def _cont_wrapper(*a: P.args, **k: P.kwargs) -> T:
+        with handler({_get_args: lambda: (a, k)}):
+            return fn(*a, **k)
+
+    return _cont_wrapper
+
+
+def _set_prompt(
+    prompt: Operation[P, T], cont: Callable[P, T], body: Callable[P, T]
+) -> Callable[P, T]:
+    from effectful.ops.semantics import handler
+
+    @functools.wraps(body)
+    def bound_body(*a: P.args, **k: P.kwargs) -> T:
+        next_cont = get_interpretation().get(prompt, prompt.__default_rule__)
+        with handler({prompt: handler({prompt: next_cont})(cont)}):
+            return body(*a, **k)
+
+    return bound_body

effectful/ops/semantics.py

@@ -0,0 +1,329 @@
+import contextlib
+import functools
+from typing import Any, Callable, Optional, Set, Type, TypeVar
+
+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")
+
+
+@defop  # type: ignore
+def apply(
+    intp: Interpretation[S, T], op: Operation[P, S], *args: P.args, **kwargs: P.kwargs
+) -> T:
+    """Apply ``op`` to ``args``, ``kwargs`` in interpretation ``intp``.
+
+    Handling :func:`apply` changes the evaluation strategy of terms.
+
+    **Example usage**:
+
+    >>> @defop
+    ... def add(x: int, y: int) -> int:
+    ...     return x + y
+    >>> @defop
+    ... def mul(x: int, y: int) -> int:
+    ...     return x * y
+
+    ``add`` and ``mul`` have default rules, so this term evaluates:
+
+    >>> mul(add(1, 2), 3)
+    9
+
+    By installing an :func:`apply` handler, we capture the term instead:
+
+    >>> with handler({apply: lambda _, op, *args, **kwargs: op.__free_rule__(*args, **kwargs) }):
+    ...     term = mul(add(1, 2), 3)
+    >>> term
+    mul(add(1, 2), 3)
+
+    """
+    if op in intp:
+        return intp[op](*args, **kwargs)
+    elif apply in intp:
+        return intp[apply](intp, 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.
+
+    :func:`fwd` should only be called in the context of a handler.
+
+    :param args: Positional arguments.
+    :param kwargs: Keyword arguments.
+
+    If no positional or keyword arguments are provided, :func:`fwd` will forward
+    the current arguments to the next handler.
+
+    """
+    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]:
+    """The coproduct of two interpretations handles any effect that is handled
+    by either. If both interpretations handle an effect, ``intp2`` takes
+    precedence.
+
+    Handlers in ``intp2`` that override a handler in ``intp`` may call the
+    overridden handler using :func:`fwd`. This allows handlers to be written
+    that extend or wrap other handlers.
+
+    **Example usage**:
+
+    The ``message`` effect produces a welcome message using two helper effects:
+    ``greeting`` and ``name``. By handling these helper effects, we can customize the
+    message.
+
+    >>> message, greeting, name = defop(str), defop(str), defop(str)
+    >>> i1 = {message: lambda: f"{greeting()} {name()}!", greeting: lambda: "Hi"}
+    >>> i2 = {name: lambda: "Jack"}
+
+    The coproduct of ``i1`` and ``i2`` handles all three effects.
+
+    >>> i3 = coproduct(i1, i2)
+    >>> with handler(i3):
+    ...     print(f'{message()}')
+    Hi Jack!
+
+    We can delegate to an enclosing handler by calling :func:`fwd`. Here we
+    override the ``name`` handler to format the name differently.
+
+    >>> i4 = coproduct(i3, {name: lambda: f'*{fwd()}*'})
+    >>> with handler(i4):
+    ...     print(f'{message()}')
+    Hi *Jack*!
+
+    .. note::
+
+      :func:`coproduct` allows effects to be overridden in a pervasive way, but
+      this is not always desirable. In particular, an interpretation with
+      handlers that call "internal" private effects may be broken if coproducted
+      with an interpretation that handles those effects. It is dangerous to take
+      the coproduct of arbitrary interpretations. For an alternate form of
+      interpretation composition, see :func:`product`.
+
+    """
+    from effectful.internals.runtime import (
+        _get_args,
+        _restore_args,
+        _save_args,
+        _set_prompt,
+    )
+
+    res = dict(intp)
+    for op, i2 in intp2.items():
+        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
+
+            # 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))
+
+    return res
+
+
+def product(
+    intp: Interpretation[S, T], intp2: Interpretation[S, T]
+) -> Interpretation[S, T]:
+    """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,
+    ``intp`` is isolated from ``intp2``.
+
+    **Example usage**:
+
+    In this example, ``i1`` has a ``param`` effect that defines some hyperparameter and
+    an effect ``f1`` that uses it. ``i2`` redefines ``param`` and uses it in a new effect
+    ``f2``, which calls ``f1``.
+
+    >>> param, f1, f2 = defop(int), defop(dict), defop(dict)
+    >>> i1 = {param: lambda: 1, f1: lambda: {'inner': param()}}
+    >>> i2 = {param: lambda: 2, f2: lambda: f1() | {'outer': param()}}
+
+    Using :func:`product`, ``i2``'s override of ``param`` is not visible to ``i1``.
+
+    >>> with handler(product(i1, i2)):
+    ...     print(f2())
+    {'inner': 1, 'outer': 2}
+
+    However, if we use :func:`coproduct`, ``i1`` is not isolated from ``i2``.
+
+    >>> with handler(coproduct(i1, i2)):
+    ...     print(f2())
+    {'inner': 2, 'outer': 2}
+
+    **References**
+
+    [1] Ahman, D., & Bauer, A. (2020, April). Runners in action. In European
+    Symposium on Programming (pp. 29-55). Cham: Springer International
+    Publishing.
+
+    """
+    if any(op in intp for op in intp2):  # alpha-rename
+        renaming = {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:
+        refls2 = {op: op.__default_rule__ for op in intp2}
+        intp_ = coproduct({}, {op: runner(refls2)(intp[op]) for op in intp})
+        return {op: runner(intp_)(intp2[op]) for op in intp2}
+
+
+@contextlib.contextmanager
+def runner(intp: Interpretation[S, T]):
+    """Install an interpretation by taking a product with the current
+    interpretation.
+
+    """
+    from effectful.internals.runtime import get_interpretation, interpreter
+
+    @interpreter(get_interpretation())
+    def _reapply(_, op: Operation[P, S], *args: P.args, **kwargs: P.kwargs):
+        return op(*args, **kwargs)
+
+    with interpreter({apply: _reapply, **intp}):
+        yield intp
+
+
+@contextlib.contextmanager
+def handler(intp: Interpretation[S, T]):
+    """Install an interpretation by taking a coproduct with the current
+    interpretation.
+
+    """
+    from effectful.internals.runtime import get_interpretation, interpreter
+
+    with interpreter(coproduct(get_interpretation(), intp)):
+        yield intp
+
+
+def evaluate(expr: Expr[T], *, intp: Optional[Interpretation[S, T]] = None) -> Expr[T]:
+    """Evaluate expression ``expr`` using interpretation ``intp``. If no
+    interpretation is provided, uses the current interpretation.
+
+    :param expr: The expression to evaluate.
+    :param intp: Optional interpretation for evaluating ``expr``.
+
+    **Example usage**:
+
+    >>> @defop
+    ... def add(x: int, y: int) -> int:
+    ...     raise NoDefaultRule
+    >>> expr = add(1, add(2, 3))
+    >>> 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()
+
+    expr = defterm(expr) if not isinstance(expr, Term) else 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)  # type: ignore
+    elif tree.is_nested(expr):
+        return tree.map_structure(functools.partial(evaluate, intp=intp), expr)
+    else:
+        return expr
+
+
+def typeof(term: Expr[T]) -> Type[T]:
+    """Return the type of an expression.
+
+    **Example usage**:
+
+    Type signatures are used to infer the types of expressions.
+
+    >>> @defop
+    ... def cmp(x: int, y: int) -> bool:
+    ...     raise NoDefaultRule
+    >>> 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
+    >>> 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
+
+
+def fvsof(term: Expr[S]) -> 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}
+
+    """
+    from effectful.internals.runtime import interpreter
+
+    _fvs: Set[Operation] = set()
+
+    def _update_fvs(_, op, *args, **kwargs):
+        _fvs.add(op)
+        for bound_var in op.__fvs_rule__(*args, **kwargs):
+            if bound_var in _fvs:
+                _fvs.remove(bound_var)
+
+    with interpreter({apply: _update_fvs}):
+        evaluate(term)
+
+    return _fvs