From 89886558ce9a05b003c0651c223e0b6b73f4aa38 Mon Sep 17 00:00:00 2001 From: Conrad Date: Mon, 13 Jul 2026 15:13:01 -0400 Subject: [PATCH 1/2] fix: Support contextlib stacks for noreentry-guarded context managers Pushing a guarded context manager onto an ExitStack or AsyncExitStack raised TypeError, the error the decorator emits for bare functions. contextlib looks special methods up on the type and calls them unbound, so type(cm).__enter__(cm) reaches the descriptor itself rather than a bound wrapper, and __call__ rejected it. WorkerPool was affected: a pool could not be composed onto a stack alongside other resources. Record the owning class via __set_name__ and treat an owned descriptor invoked with an instance of its owner as the bound call it actually is, delegating through __get__ so the guard, the coroutine-ness, and the wrapper cache all still apply. Class-level access keeps returning the descriptor, so introspection is unchanged, and a bare function still has no owner and is still rejected. Teardown works the same way: contextlib registers exit as a method of the type, so a guarded __exit__ receives its exception triple through the same path. --- wool/src/wool/utilities/noreentry.py | 104 +++++++++++++++++++-------- 1 file changed, 75 insertions(+), 29 deletions(-) diff --git a/wool/src/wool/utilities/noreentry.py b/wool/src/wool/utilities/noreentry.py index 50a6867a..42f4551f 100644 --- a/wool/src/wool/utilities/noreentry.py +++ b/wool/src/wool/utilities/noreentry.py @@ -6,7 +6,6 @@ import sys import weakref from typing import Callable -from typing import Never from typing import ParamSpec from typing import TypeVar @@ -20,19 +19,51 @@ class _Token: pass -class NoReentryBoundMethod: - """Descriptor implementing single-use guard for bound methods. - - On the first call the decorated method executes normally. Any - subsequent call raises :class:`RuntimeError`. - - Guard state uses a per-instance token stored on the instance under - ``__noreentry_token__``. The token is unique, hashable, and tied to the - instance's lifetime. The descriptor tracks tokens in a WeakSet to auto-clean - when instances are garbage collected. - - Works with both synchronous and asynchronous methods. Only supports - bound methods; using @noreentry on bare functions raises TypeError. +class NoReentryDescriptor: + """Guard a method against a second invocation on the same instance. + + A descriptor that wraps one method of one class. It exists for the + once-per-instance operations — teardown, shutdown, close — whose second + execution is a latent bug rather than a no-op, and turns that bug into an + immediate, attributable failure. + + The guard is per instance, not per class: each instance of the owning class + may invoke the method once. The first call executes the method and returns + its result; every later call on the same instance raises `RuntimeError`. + Guard state lives on the instance under ``__noreentry_token__`` and is + discarded with it, so a fresh instance is unguarded and an instance that is + garbage collected leaves nothing behind. + + Because the guard keys on the receiver, the method must be reached with one: + it may be called on an instance, or unbound off the owning class with the + instance as the first argument, i.e., the form `contextlib` uses to enter a + context manager (``type(cm).__enter__(cm)``). The latter reaches the + descriptor rather than a bound wrapper, so a guarded ``__enter__`` or + ``__aenter__`` composes with `contextlib.ExitStack` and + `contextlib.AsyncExitStack`. Synchronous and asynchronous methods are both + supported. A receiver of any other type, or none at all, raises `TypeError`, + as does decorating a bare function, which has no receiver to key on. + + .. rubric:: Implementation notes + + Owner tracking is the reason this is a descriptor rather than a plain + function decorator. A decorator that leans on Python's descriptor protocol + to do the binding never sees the unbound call form at all — the arguments + arrive as an ordinary positional tuple — so it cannot tell an unbound call + on the owning instance from a call whose first argument merely happens to be + an object, and would silently run the method against a foreign receiver. + Holding the owner lets that case be rejected. + + ``__set_name__`` supplies the discriminator: the interpreter invokes it only + for a descriptor assigned in a class body, so ``_owner`` stays ``None`` for a + bare function. A ``None`` owner therefore means "no receiver exists" and a + set owner means "a receiver is required", which is the distinction ``__call__`` + dispatches on. + + An owned descriptor called with its own instance first is a bound call in + disguise: rebinding through ``__get__`` and re-dispatching routes it through + the same wrapper and the same guard state as an instance call, so the two + call forms cannot diverge. """ def __init__(self, fn, /): @@ -43,10 +74,34 @@ def __init__(self, fn, /): else: self._is_coroutine = asyncio.coroutines._is_coroutine # type: ignore[attr-defined] self._fn = fn + self._owner = None self._invocations = weakref.WeakSet() - def __call__(self, *args, **kwargs) -> Never: - raise TypeError("@noreentry only decorates methods, not bare functions") + def __set_name__(self, owner, name): + """Record the class the descriptor is defined on.""" + self._owner = owner + + def __call__(self, *args, **kwargs): + """Dispatch an unbound call to the receiver's bound wrapper. + + This is the path `contextlib` takes when it enters a context manager by + invoking the special method off the type (``type(cm).__enter__(cm)``). + See the class docstring for the guard semantics and the supported call + forms. + + :raises TypeError: + If the decorated object is a bare function, or if the first + positional argument is not an instance of the owning class. + """ + if self._owner is None: + raise TypeError("@noreentry only decorates methods, not bare functions") + if not args or not isinstance(args[0], self._owner): + raise TypeError( + f"'{self._fn.__qualname__}' must be called with an instance of " + f"'{self._owner.__name__}' as the first argument" + ) + obj, *rest = args + return self.__get__(obj, self._owner)(*rest, **kwargs) def __get__(self, obj, objtype=None): if obj is None: @@ -95,19 +150,10 @@ def _guard(self, obj): def noreentry(fn: Callable[P, R]) -> Callable[P, R]: """Mark a method as single-use. - On the first call the decorated method executes normally. Any - subsequent call raises :class:`RuntimeError`. - - Guard state uses a per-instance token stored on the instance under - ``__noreentry_token__``. The token is unique, hashable, and tied to the - instance's lifetime. The descriptor tracks tokens in a :class:`WeakSet` - to auto-clean when instances are garbage collected. - - Works with both synchronous and asynchronous methods. Only supports - bound methods; using @noreentry on bare functions raises :class:`TypeError` - on the first invocation. + Returns a `NoReentryDescriptor`; see that class for the guard semantics and + the supported call forms. :param fn: - The bound instance or class method to decorate. + The instance method to decorate. """ - return NoReentryBoundMethod(fn) + return NoReentryDescriptor(fn) From 63ad2718973a383b5247bfbba9f282ea5b5e6666 Mon Sep 17 00:00:00 2001 From: Conrad Date: Mon, 13 Jul 2026 15:13:14 -0400 Subject: [PATCH 2/2] test: Cover contextlib composition and mixed call forms for noreentry Pin the reported failure at both the utility and the WorkerPool level: a guarded manager enters through ExitStack and AsyncExitStack, and the guard still rejects a second entry there. Guard the argument-forwarding half of the contract, which nothing exercised. contextlib registers teardown as a method of the type, so a guarded exit is routed through the descriptor with its exception triple as trailing arguments; every prior test entered with no arguments at all, so an implementation that dropped them silently passed. Cover the exception details reaching a guarded exit through both stacks, and argument forwarding on a direct unbound call. Pin the guard as shared across the bound and unbound call forms, by example and by a property over arbitrary interleavings of the two, so the widened dispatch cannot become a way around single-use semantics. Cover a subclass entering through a stack, since dispatch admits any instance of the owning class rather than only exact instances, and the plain statement form, which had no coverage of its own. Rename the bare-function test to say what it asserts: it calls an owned descriptor with no instance, which is a distinct rejection branch from a genuinely unowned function. Drop the unbound-access test, whose only assertion could not fail. --- wool/tests/runtime/worker/test_pool.py | 53 +++ wool/tests/utilities/test_noreentry.py | 469 ++++++++++++++++++++++++- 2 files changed, 508 insertions(+), 14 deletions(-) diff --git a/wool/tests/runtime/worker/test_pool.py b/wool/tests/runtime/worker/test_pool.py index 9609a9e5..f48ad27b 100644 --- a/wool/tests/runtime/worker/test_pool.py +++ b/wool/tests/runtime/worker/test_pool.py @@ -10,6 +10,7 @@ import os import time import uuid +from contextlib import AsyncExitStack from contextlib import contextmanager from functools import partial from types import MappingProxyType @@ -1190,6 +1191,58 @@ async def test___aenter___lifecycle(self, mock_worker_factory): # Assert: Context manager returned pool and lifecycle completed assert pool_instance is not None + @pytest.mark.asyncio + async def test___aenter___should_return_pool_when_pushed_onto_async_exit_stack( + self, + mock_shared_memory, + mock_worker_proxy, + mock_local_worker, + ): + """Test entering a pool through an AsyncExitStack. + + Given: + A WorkerPool + When: + The pool is pushed onto a contextlib.AsyncExitStack + Then: + It should return the pool itself. + """ + # Arrange + pool = WorkerPool(spawn=1) + + # Act + async with AsyncExitStack() as stack: + entered = await stack.enter_async_context(pool) + + # Assert + assert entered is pool + + @pytest.mark.asyncio + async def test___aexit___should_tear_down_pool_when_async_exit_stack_unwinds( + self, + mock_shared_memory, + mock_worker_proxy, + mock_local_worker, + ): + """Test tearing a pool down through an AsyncExitStack. + + Given: + A WorkerPool pushed onto a contextlib.AsyncExitStack + When: + The stack unwinds + Then: + It should exit the pool's worker proxy context. + """ + # Arrange + pool = WorkerPool(spawn=1) + + # Act + async with AsyncExitStack() as stack: + await stack.enter_async_context(pool) + + # Assert + mock_worker_proxy.__aexit__.assert_awaited() + @pytest.mark.asyncio async def test___aexit___cleanup_on_error(self, mock_worker_factory): """Test cleanup still occurs and exception propagates correctly. diff --git a/wool/tests/utilities/test_noreentry.py b/wool/tests/utilities/test_noreentry.py index 1e0b08c8..52c03be2 100644 --- a/wool/tests/utilities/test_noreentry.py +++ b/wool/tests/utilities/test_noreentry.py @@ -1,8 +1,13 @@ from __future__ import annotations import asyncio +from contextlib import AsyncExitStack +from contextlib import ExitStack import pytest +from hypothesis import given +from hypothesis import settings +from hypothesis import strategies as st from wool.utilities.noreentry import noreentry @@ -24,6 +29,66 @@ async def run(self): return "ok" +class _SyncContextDummy: + """Sync context manager with both __enter__ and __exit__ guarded. + + Records the exception type handed to __exit__ so the arguments + forwarded through the descriptor can be observed. + """ + + def __init__(self): + self.exited_with = "not-exited" + + @noreentry + def __enter__(self): + return self + + @noreentry + def __exit__(self, exc_type, exc, tb): + self.exited_with = exc_type + return False + + +class _AsyncContextDummy: + """Async context manager with both __aenter__ and __aexit__ guarded. + + Records the exception type handed to __aexit__ so the arguments + forwarded through the descriptor can be observed. + """ + + def __init__(self): + self.exited_with = "not-exited" + + @noreentry + async def __aenter__(self): + return self + + @noreentry + async def __aexit__(self, exc_type, exc, tb): + self.exited_with = exc_type + return False + + +class _SubclassedContextDummy(_SyncContextDummy): + """Subclass inheriting a @noreentry-guarded __enter__ from its base.""" + + +class _ArgumentDummy: + """Class with a @noreentry method taking positional and keyword args.""" + + @noreentry + def run(self, first, *, second=None): + return first, second + + +class _PayloadDummy: + """Class with a @noreentry method echoing back its whole payload.""" + + @noreentry + def run(self, *args, **kwargs): + return args, kwargs + + class _MultiMethodDummy: """Class with multiple @noreentry methods.""" @@ -36,6 +101,28 @@ def second(self): return "second" +@noreentry +def _bare_function(): + """Module-level function decorated with @noreentry, owned by no class.""" + return "ok" + + +_payloads = st.one_of( + st.none(), + st.booleans(), + st.integers(), + st.text(), + st.binary(), + st.lists(st.integers()), +) + +_keywords = st.text( + alphabet=st.characters(min_codepoint=97, max_codepoint=122), + min_size=1, + max_size=8, +).filter(lambda name: name != "self") + + class TestNoReentry: """Tests for the noreentry decorator.""" @@ -203,29 +290,53 @@ def test_noreentry_multiple_methods_independent(self): # Assert assert result == "second" - def test_noreentry_unbound_access_returns_descriptor(self): - """Test accessing decorated method via class returns descriptor. + def test_noreentry_should_raise_type_error_when_called_unbound_without_instance( + self, + ): + """Test an unbound call with no instance is rejected. Given: - A class with a @noreentry method. + A @noreentry method accessed unbound through its class When: - The method is accessed through the class (unbound). + It is called with no arguments, so there is no instance to + bind Then: - It should return the descriptor itself. + It should raise TypeError. """ - # Act + # Arrange unbound = _SyncDummy.run - # Assert - assert unbound is not None + # Act & assert + with pytest.raises( + TypeError, match="must be called with an instance of '_SyncDummy'" + ): + unbound() + + def test_noreentry_should_raise_type_error_when_bare_function_is_called(self): + """Test a decorated module-level function rejects invocation. + + Given: + A @noreentry-decorated module-level function, owned by no + class + When: + The function is called + Then: + It should raise TypeError. + """ + # Act & assert + with pytest.raises( + TypeError, match="only decorates methods, not bare functions" + ): + _bare_function() - def test_noreentry_bare_function_raises_error(self): - """Test decorator rejects application to bare functions. + def test_noreentry_should_raise_type_error_when_called_with_foreign_instance(self): + """Test an unbound call rejects an instance it does not own. Given: - A @noreentry-decorated function called without a bound instance. + A @noreentry method accessed unbound through its class When: - The descriptor is called directly. + It is called with an object that is not an instance of the + class that owns it Then: It should raise TypeError. """ @@ -234,6 +345,336 @@ def test_noreentry_bare_function_raises_error(self): # Act & assert with pytest.raises( - TypeError, match="only decorates methods, not bare functions" + TypeError, match="must be called with an instance of '_SyncDummy'" ): - unbound() + unbound(_MultiMethodDummy()) + + def test_noreentry_should_enter_context_when_guarded_enter_pushed_onto_stack(self): + """Test a guarded sync manager enters through an ExitStack. + + Given: + A sync context manager whose __enter__ is decorated with + @noreentry + When: + The manager is pushed onto a contextlib.ExitStack + Then: + It should enter successfully and yield the manager. + """ + # Arrange + manager = _SyncContextDummy() + + # Act + with ExitStack() as stack: + entered = stack.enter_context(manager) + + # Assert + assert entered is manager + + def test_noreentry_should_raise_when_guarded_enter_pushed_onto_stack_twice(self): + """Test the guard still fires for a manager entered on a stack. + + Given: + A sync context manager already entered through an ExitStack + When: + The same manager is pushed onto the stack a second time + Then: + It should raise RuntimeError. + """ + # Arrange + manager = _SyncContextDummy() + + # Act & assert + with ExitStack() as stack: + stack.enter_context(manager) + with pytest.raises(RuntimeError, match="cannot be invoked more than once"): + stack.enter_context(manager) + + @pytest.mark.asyncio + async def test_noreentry_should_enter_context_when_guarded_aenter_pushed_onto_stack( + self, + ): + """Test a guarded async manager enters through an AsyncExitStack. + + Given: + An async context manager whose __aenter__ is decorated with + @noreentry + When: + The manager is pushed onto a contextlib.AsyncExitStack + Then: + It should enter successfully and yield the manager. + """ + # Arrange + manager = _AsyncContextDummy() + + # Act + async with AsyncExitStack() as stack: + entered = await stack.enter_async_context(manager) + + # Assert + assert entered is manager + + @pytest.mark.asyncio + async def test_noreentry_should_raise_when_guarded_aenter_pushed_onto_stack_twice( + self, + ): + """Test the guard still fires for an async manager on a stack. + + Given: + An async context manager already entered through an + AsyncExitStack + When: + The same manager is pushed onto the stack a second time + Then: + It should raise RuntimeError. + """ + # Arrange + manager = _AsyncContextDummy() + + # Act & assert + async with AsyncExitStack() as stack: + await stack.enter_async_context(manager) + with pytest.raises(RuntimeError, match="cannot be invoked more than once"): + await stack.enter_async_context(manager) + + def test_noreentry_should_enter_context_when_guarded_enter_used_in_with(self): + """Test a guarded manager works in a plain with statement. + + Given: + A sync context manager whose __enter__ is decorated with + @noreentry + When: + The manager is used in a plain with statement + Then: + It should enter and yield the manager. + """ + # Arrange + manager = _SyncContextDummy() + + # Act + with manager as entered: + # Assert + assert entered is manager + + @pytest.mark.asyncio + async def test_noreentry_should_enter_context_when_guarded_aenter_used_in_with(self): + """Test a guarded manager works in a plain async with statement. + + Given: + An async context manager whose __aenter__ is decorated with + @noreentry + When: + The manager is used in a plain async with statement + Then: + It should enter and yield the manager. + """ + # Arrange + manager = _AsyncContextDummy() + + # Act + async with manager as entered: + # Assert + assert entered is manager + + def test_noreentry_should_forward_exception_details_when_exit_stack_unwinds(self): + """Test a guarded __exit__ receives the exception forwarded to it. + + Given: + A context manager whose __enter__ and __exit__ are both + decorated with @noreentry + When: + An exception is raised inside an ExitStack body holding the + manager + Then: + It should forward the exception details to __exit__ and let + the exception propagate. + """ + # Arrange + manager = _SyncContextDummy() + + # Act + with pytest.raises(ValueError, match="boom"): + with ExitStack() as stack: + stack.enter_context(manager) + raise ValueError("boom") + + # Assert — contextlib invokes the guarded __exit__ unbound, so + # the exception triple must survive the descriptor + assert manager.exited_with is ValueError + + @pytest.mark.asyncio + async def test_noreentry_should_forward_exception_details_when_async_stack_unwinds( + self, + ): + """Test a guarded __aexit__ receives the exception forwarded to it. + + Given: + An async context manager whose __aenter__ and __aexit__ are + both decorated with @noreentry + When: + An exception is raised inside an AsyncExitStack body holding + the manager + Then: + It should forward the exception details to __aexit__ and let + the exception propagate. + """ + # Arrange + manager = _AsyncContextDummy() + + # Act + with pytest.raises(ValueError, match="boom"): + async with AsyncExitStack() as stack: + await stack.enter_async_context(manager) + raise ValueError("boom") + + # Assert + assert manager.exited_with is ValueError + + def test_noreentry_should_raise_when_entered_by_statement_then_exit_stack(self): + """Test the guard is shared across bound and unbound call forms. + + Given: + A guarded context manager already entered with a plain with + statement + When: + The same manager is pushed onto an ExitStack, which enters it + through the unbound call form + Then: + It should raise RuntimeError, the guard being shared across + both forms. + """ + # Arrange + manager = _SyncContextDummy() + with manager: + pass + + # Act & assert + with ExitStack() as stack: + with pytest.raises(RuntimeError, match="cannot be invoked more than once"): + stack.enter_context(manager) + + @pytest.mark.asyncio + async def test_noreentry_should_raise_when_entered_by_async_with_then_stack(self): + """Test the async guard is shared across bound and unbound forms. + + Given: + A guarded async context manager already entered with a plain + async with statement + When: + The same manager is pushed onto an AsyncExitStack, which + enters it through the unbound call form + Then: + It should raise RuntimeError, the guard being shared across + both forms. + """ + # Arrange + manager = _AsyncContextDummy() + async with manager: + pass + + # Act & assert + async with AsyncExitStack() as stack: + with pytest.raises(RuntimeError, match="cannot be invoked more than once"): + await stack.enter_async_context(manager) + + def test_noreentry_should_enter_context_when_subclass_pushed_onto_exit_stack(self): + """Test a subclass inherits a guarded __enter__ usably. + + Given: + A subclass of a context manager whose __enter__ is decorated + with @noreentry on the base class + When: + An instance of the subclass is pushed onto an ExitStack + Then: + It should enter successfully and yield the manager. + """ + # Arrange + manager = _SubclassedContextDummy() + + # Act + with ExitStack() as stack: + entered = stack.enter_context(manager) + + # Assert + assert entered is manager + + def test_noreentry_should_forward_arguments_when_called_unbound(self): + """Test an unbound call forwards arguments past the instance. + + Given: + A class with a @noreentry method taking a positional and a + keyword-only argument + When: + The method is called unbound through its class with the + instance first, followed by both arguments + Then: + It should return the result computed from both forwarded + arguments. + """ + # Arrange + obj = _ArgumentDummy() + + # Act + result = _ArgumentDummy.run(obj, "first", second="second") + + # Assert + assert result == ("first", "second") + + @given( + args=st.lists(_payloads, max_size=4), + kwargs=st.dictionaries(_keywords, _payloads, max_size=4), + ) + @settings(max_examples=50, deadline=None) + def test_noreentry_should_forward_arbitrary_payload_when_called_unbound( + self, args, kwargs + ): + """Test an unbound call forwards any payload past the instance. + + Given: + Any tuple of positional arguments and any mapping of keyword + arguments + When: + A guarded echo method is called unbound through its class + with a fresh instance first, followed by the payload + Then: + It should return exactly the payload it was given. + """ + # Arrange — a fresh instance per example, so no example inherits + # another's exhausted guard + obj = _PayloadDummy() + + # Act + result = _PayloadDummy.run(obj, *args, **kwargs) + + # Assert + assert result == (tuple(args), kwargs) + + @given(forms=st.lists(st.sampled_from(["bound", "unbound"]), min_size=1, max_size=6)) + @settings(max_examples=25, deadline=None) + def test_noreentry_should_admit_only_the_first_call_when_call_forms_interleave( + self, forms + ): + """Test at most one invocation survives any mix of call forms. + + Given: + Any non-empty sequence of invocation forms drawn from the + bound and unbound call shapes + When: + The sequence is applied to a single fresh instance + Then: + It should let exactly the first invocation through and raise + RuntimeError for every later one, whatever the interleaving. + """ + # Arrange — a fresh instance per example, so no example inherits + # another's exhausted guard + obj = _SyncDummy() + + def invoke(form): + if form == "bound": + return obj.run() + return _SyncDummy.run(obj) + + # Act & assert + assert invoke(forms[0]) == "ok" + for form in forms[1:]: + with pytest.raises(RuntimeError, match="cannot be invoked more than once"): + invoke(form)