Sync main with release

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)

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.