From c63522c4d4f4e0f3f223b574b6d80a89dc76ba3c Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 31 Aug 2026 16:08:21 -0400 Subject: [PATCH 1/2] forwarding proceeds through apply to reach default rule --- effectful/ops/types.py | 29 ++++++--- tests/test_ops_semantics.py | 120 +++++++++++++++++++++++++++++++++++- 2 files changed, 140 insertions(+), 9 deletions(-) diff --git a/effectful/ops/types.py b/effectful/ops/types.py index 6d9b99c19..6559b1221 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -510,10 +510,15 @@ def _instance_op(instance, *args, **kwargs): return self @functools.cached_property - def _default_rule_with_args(self): + def _next_rule_with_args(self): from effectful.internals.runtime import _restore_args - return _restore_args(self.__default_rule__) + rule = ( + self.__default_rule__ + if isinstance(self, ApplyOperation) + else functools.partial(self.__apply__, self) + ) + return _restore_args(rule) def __call__(self, *args: Q.args, **kwargs: Q.kwargs) -> V: from effectful.internals.runtime import get_interpretation @@ -523,9 +528,9 @@ def __call__(self, *args: Q.args, **kwargs: Q.kwargs) -> V: self_handler = intp.get(self) if self_handler is not None: - # ensure that fwd is bound to the default rule. if this handler has - # a bound fwd, it will override this binding - fwd_intp = typing.cast(Interpretation, {fwd: self._default_rule_with_args}) + # Operation handlers forward through apply before reaching the default. + # Apply handlers forward to their own generated/default implementation. + fwd_intp = typing.cast(Interpretation, {fwd: self._next_rule_with_args}) with handler(fwd_intp): return self_handler(*args, **kwargs) elif args and isinstance(args[0], Operation) and self is args[0].__apply__: @@ -534,11 +539,15 @@ def __call__(self, *args: Q.args, **kwargs: Q.kwargs) -> V: else: return self.__apply__(self, *args, **kwargs) - def __init_subclass__(cls, **kwargs) -> None: + def __init_subclass__(cls, *, _generate_apply: bool = True, **kwargs) -> None: + super().__init_subclass__(**kwargs) + if not _generate_apply: + return + assert "__apply__" not in cls.__dict__ or cls is Operation, ( "Cannot manually override apply" ) - assert isinstance(cls.__apply__, Operation) + assert isinstance(cls.__apply__, ApplyOperation) cls.__apply__ = cls.__apply__.define( staticmethod( @@ -552,6 +561,10 @@ def __init_subclass__(cls, **kwargs) -> None: ) +class ApplyOperation[**Q, V](Operation[Q, V], _generate_apply=False): + """An operation that implements application for an Operation subclass.""" + + def __apply__[**A, B](op: Operation[A, B], *args: A.args, **kwargs: A.kwargs) -> B: """Apply ``op`` to ``args``, ``kwargs`` in interpretation ``intp``. @@ -584,7 +597,7 @@ def __apply__[**A, B](op: Operation[A, B], *args: A.args, **kwargs: A.kwargs) -> return op.__default_rule__(*args, **kwargs) # type: ignore[return-value] -Operation.__apply__ = Operation.define(staticmethod(__apply__)) +Operation.__apply__ = ApplyOperation.define(staticmethod(__apply__)) del __apply__ diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index 17c542561..9ccba86e0 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -23,7 +23,13 @@ defop, implements, ) -from effectful.ops.types import Interpretation, NotHandled, Operation, Term +from effectful.ops.types import ( + ApplyOperation, + Interpretation, + NotHandled, + Operation, + Term, +) logger = logging.getLogger(__name__) @@ -319,6 +325,118 @@ def plus_1_fwd(x): assert plus_1(1) == 2 +def test_fwd_from_operation_handler_to_apply_handler(): + @Operation.define + def f(x: int) -> int: + return x + 1 + + calls = [] + + def f_handler(x): + calls.append(("f", x)) + return fwd() + + def apply_handler(op, *args, **kwargs): + calls.append(("apply", op, args, kwargs)) + return fwd() + + with handler({apply: apply_handler}), handler({f: f_handler}): + assert f(1) == 2 + + assert calls == [("f", 1), ("apply", f, (1,), {})] + + +def test_fwd_from_operation_handler_to_apply_handler_with_replacement_args(): + @Operation.define + def f(x: int) -> int: + return x + + apply_args = [] + + def apply_handler(op, *args, **kwargs): + apply_args.append((op, args, kwargs)) + return fwd(op, *args, **kwargs) + + with handler({apply: apply_handler}), handler({f: lambda x: fwd(x + 1)}): + assert f(1) == 2 + + assert apply_args == [(f, (2,), {})] + + +def test_fwd_through_apply_handlers_is_associative(): + calls = [] + + @Operation.define + def f() -> int: + calls.append("default") + return 1 + + def forwarding(name): + def impl(*args, **kwargs): + calls.append(name) + return fwd() + + return impl + + h0 = {apply: forwarding("left apply")} + h1 = {f: forwarding("exact")} + h2 = {apply: forwarding("right apply")} + expected = ["exact", "right apply", "left apply", "default"] + + for intp in ( + coproduct(coproduct(h0, h1), h2), + coproduct(h0, coproduct(h1, h2)), + ): + calls.clear() + with handler(intp): + assert f() == 1 + assert calls == expected + + calls.clear() + with handler(h0), handler(h1), handler(h2): + assert f() == 1 + assert calls == expected + + +def test_fwd_through_apply_operation_subtypes(): + calls = [] + + class BaseOperation(Operation): + pass + + class DerivedOperation(BaseOperation): + pass + + @DerivedOperation.define + def f(x: int) -> int: + calls.append("default") + return x + 1 + + def forwarding(name): + def impl(*args, **kwargs): + calls.append(name) + return fwd() + + return impl + + assert isinstance(Operation.__apply__, ApplyOperation) + assert isinstance(BaseOperation.__apply__, ApplyOperation) + assert isinstance(DerivedOperation.__apply__, ApplyOperation) + assert not isinstance(f, ApplyOperation) + + with handler( + { + Operation.__apply__: forwarding("apply"), + BaseOperation.__apply__: forwarding("base apply"), + DerivedOperation.__apply__: forwarding("derived apply"), + f: forwarding("exact"), + } + ): + assert f(1) == 2 + + assert calls == ["exact", "derived apply", "base apply", "apply", "default"] + + @pytest.mark.parametrize("op,args", OPERATION_CASES) @pytest.mark.parametrize("n1", N_CASES) @pytest.mark.parametrize("n2", N_CASES) From a096cd8fff0e09c3fdb033aeeafd12a8a57c1783 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 31 Aug 2026 16:19:35 -0400 Subject: [PATCH 2/2] simplify --- effectful/ops/types.py | 7 ++----- tests/test_ops_semantics.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/effectful/ops/types.py b/effectful/ops/types.py index 6559b1221..f3a830c19 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -533,8 +533,7 @@ def __call__(self, *args: Q.args, **kwargs: Q.kwargs) -> V: fwd_intp = typing.cast(Interpretation, {fwd: self._next_rule_with_args}) with handler(fwd_intp): return self_handler(*args, **kwargs) - elif args and isinstance(args[0], Operation) and self is args[0].__apply__: - # Prevent infinite recursion when calling self.apply directly + elif isinstance(self, ApplyOperation): return self.__default__(*args, **kwargs) else: return self.__apply__(self, *args, **kwargs) @@ -544,9 +543,7 @@ def __init_subclass__(cls, *, _generate_apply: bool = True, **kwargs) -> None: if not _generate_apply: return - assert "__apply__" not in cls.__dict__ or cls is Operation, ( - "Cannot manually override apply" - ) + assert "__apply__" not in cls.__dict__, "Cannot manually override apply" assert isinstance(cls.__apply__, ApplyOperation) cls.__apply__ = cls.__apply__.define( diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index 9ccba86e0..604edc36f 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -436,6 +436,19 @@ def impl(*args, **kwargs): assert calls == ["exact", "derived apply", "base apply", "apply", "default"] + # Unhandled intermediate apply operations proceed directly to their defaults, + # so the base apply handler sees the original operation exactly once. + calls.clear() + with handler( + { + Operation.__apply__: forwarding("apply"), + f: forwarding("exact"), + } + ): + assert f(1) == 2 + + assert calls == ["exact", "apply", "default"] + @pytest.mark.parametrize("op,args", OPERATION_CASES) @pytest.mark.parametrize("n1", N_CASES)