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) 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)