Add generator modifier to mark_as_sync and mark_as_async.

The calling convention markers now accept an optional generator keyword
argument. Tri-state: None (default, auto), True, False. Auto preserves
generator-ness across the sync/async flip -- an async generator marked
sync becomes a sync generator, a sync generator marked async becomes an
async generator. Explicit True or False force the corresponding bit.

Both markers now always clear CO_ITERABLE_COROUTINE (the legacy
@types.coroutine flag), which is incompatible with either assertion.

Tests added cover the generator parameter combinations, descriptor
binding, and stacking of the markers over with_signature. The marker
and signature decorators touch disjoint bits of co_flags so they
compose cleanly in either stacking order.

Docs in bundled.rst extended with a generator subsection under the
markers section and a new subsection under with_signature explaining
how to combine the two.

Author: GrahamDumpleton

docs/bundled.rst

@@ -468,6 +468,60 @@ changes what ``iscoroutinefunction()`` reports. The marker wrappers are for
 annotating stacks whose effective convention has already been established by
 other decorators.
 
+Marking generator convention
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Both ``mark_as_sync`` and ``mark_as_async`` accept an optional
+``generator`` keyword to control the reported generator-ness of the
+wrapper. This is the modifier on top of the primary sync/async axis, and
+lets the markers address all four realistic callable kinds: plain
+function, sync generator, coroutine function, and async generator.
+
+The parameter is tri-state:
+
+- ``None`` (default): auto. Preserve generator-ness from the input. For
+  ``mark_as_sync``, an async generator input becomes a sync generator;
+  other inputs keep their ``CO_GENERATOR`` bit as-is. For
+  ``mark_as_async``, any generator input (sync or async) becomes an
+  async generator; non-generator input becomes a coroutine function.
+- ``True``: force generator reporting on. ``mark_as_sync(generator=True)``
+  sets ``CO_GENERATOR``; ``mark_as_async(generator=True)`` sets
+  ``CO_ASYNC_GENERATOR`` (and clears ``CO_COROUTINE`` since the two
+  are mutually exclusive at the CPython code-object level).
+- ``False``: force generator reporting off. ``mark_as_sync(generator=False)``
+  clears ``CO_GENERATOR``; ``mark_as_async(generator=False)`` sets
+  ``CO_COROUTINE`` and clears ``CO_ASYNC_GENERATOR`` and ``CO_GENERATOR``.
+
+::
+
+    # An upstream decorator collects items from an async generator into
+    # a list and returns it synchronously. Mark the resulting callable
+    # as a plain sync function (default auto would flip
+    # CO_ASYNC_GENERATOR into CO_GENERATOR and the wrapper would report
+    # as a sync generator, which is wrong here -- the real return is a
+    # list).
+
+    @wrapt.mark_as_sync(generator=False)
+    @collect_async_generator_to_list
+    async def stream(...):
+        yield ...
+
+::
+
+    # A sync generator is being exposed through an adapter that wraps
+    # each yielded item in an async future. Mark it as an async
+    # generator so consumers using ``async for`` see the expected
+    # introspection.
+
+    @wrapt.mark_as_async(generator=True)
+    @async_wrap_yielded_items
+    def produce(...):
+        yield ...
+
+Both markers always clear ``CO_ITERABLE_COROUTINE`` (the legacy
+``@types.coroutine`` flag), regardless of ``generator``, since that
+flag's meaning is incompatible with either marker's assertion.
+
 Bridging between conventions
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -701,3 +755,49 @@ chain.
         ...
 
     # inspect.signature(function) still reports the prototype's signature.
+
+Combining with calling-convention markers
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+``with_signature`` controls *what the arguments look like*, but it does
+not assert anything about the *calling convention* of the resulting
+wrapper -- whether it is a coroutine function, a generator, an async
+generator, or a plain sync function. Those concerns are expressed by
+``mark_as_sync`` and ``mark_as_async``, described in the "Calling
+Convention Markers and Adapters" section above. The two concerns are
+deliberately orthogonal: ``with_signature`` only modifies the
+argument-related ``co_flags`` bits (``CO_VARARGS`` and
+``CO_VARKEYWORDS``, derived from the signature), while the markers only
+modify the calling-convention bits (``CO_COROUTINE``,
+``CO_ASYNC_GENERATOR``, ``CO_GENERATOR``, ``CO_ITERABLE_COROUTINE``).
+
+Because the two decorators touch disjoint bits, they compose cleanly
+when stacked. The conventional ordering is the marker on top, the
+signature override underneath, but both orders produce the same
+result:
+
+::
+
+    @wrapt.mark_as_sync
+    @wrapt.with_signature(prototype=_prototype)
+    async def real(*args, **kwargs):
+        # inspect.iscoroutinefunction(real) -> False (from mark_as_sync)
+        # inspect.signature(real)           -> prototype's signature
+        ...
+
+When the underlying callable is a generator of some kind and the
+surrounding stack has changed that convention, the ``generator``
+keyword on the markers is the right modifier:
+
+::
+
+    @wrapt.mark_as_async(generator=True)
+    @wrapt.with_signature(prototype=_prototype)
+    def real(*args, **kwargs):
+        # inspect.isasyncgenfunction(real) -> True
+        # inspect.signature(real)          -> prototype's signature
+        yield ...
+
+Use ``with_signature`` alone when only the signature needs correcting;
+stack with a marker when the calling convention also needs to be
+asserted independently of the wrapped function's own declaration.

docs/changes.rst

@@ -78,11 +78,15 @@ their help is much appreciated.
   ``inspect.iscoroutinefunction()`` reports the intended convention,
   letting ``synchronized`` auto-select the correct sync or async wrapping
   behaviour even when an upstream decorator has changed the effective
-  calling convention. ``async_to_sync`` runs an async callable to
-  completion via ``asyncio.run()``, and ``sync_to_async`` dispatches a
-  sync callable onto the default executor via ``loop.run_in_executor()``;
-  both self-mark so they integrate with ``synchronized`` without needing
-  an additional marker decorator. The naming of ``async_to_sync`` and
+  calling convention. Both markers take an optional ``generator`` keyword
+  (tri-state: ``None`` / ``True`` / ``False``) controlling the reported
+  generator bit, so all four callable kinds (plain function, sync
+  generator, coroutine function, async generator) can be asserted.
+  ``async_to_sync`` runs an async callable to completion via
+  ``asyncio.run()``, and ``sync_to_async`` dispatches a sync callable
+  onto the default executor via ``loop.run_in_executor()``; both
+  self-mark so they integrate with ``synchronized`` without needing an
+  additional marker decorator. The naming of ``async_to_sync`` and
   ``sync_to_async`` follows the convention used by ``asgiref``. See the
   "Calling Convention Markers and Adapters" section of :doc:`bundled` for
   details.

src/wrapt/__init__.pyi

@@ -396,8 +396,18 @@ if sys.version_info >= (3, 10):
 
     # mark_as_sync(), mark_as_async(), async_to_sync(), sync_to_async()
 
-    def mark_as_sync(wrapped: Callable[P, R]) -> Callable[P, R]: ...
-    def mark_as_async(wrapped: Callable[P, R]) -> Callable[P, R]: ...
+    @overload
+    def mark_as_sync(wrapped: Callable[P, R], /) -> Callable[P, R]: ...
+    @overload
+    def mark_as_sync(
+        *, generator: bool | None = None
+    ) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+    @overload
+    def mark_as_async(wrapped: Callable[P, R], /) -> Callable[P, R]: ...
+    @overload
+    def mark_as_async(
+        *, generator: bool | None = None
+    ) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
     def async_to_sync(wrapped: Callable[P, R]) -> Callable[P, R]: ...
     def sync_to_async(wrapped: Callable[P, R]) -> Callable[P, R]: ...
 

src/wrapt/synchronization.py

@@ -9,7 +9,13 @@
 import asyncio
 import sys
 from functools import partial
-from inspect import CO_ASYNC_GENERATOR, CO_COROUTINE, iscoroutinefunction
+from inspect import (
+    CO_ASYNC_GENERATOR,
+    CO_COROUTINE,
+    CO_GENERATOR,
+    CO_ITERABLE_COROUTINE,
+    iscoroutinefunction,
+)
 from threading import Lock, RLock
 
 from .__wrapt__ import BoundFunctionWrapper, CallableObjectProxy, FunctionWrapper
@@ -23,21 +29,40 @@
 # inner decorator that invokes an async def via asyncio.run()).
 
 
-_CO_SYNC_MASK = ~(CO_COROUTINE | CO_ASYNC_GENERATOR)
-
-
 class _SyncCodeProxy(CallableObjectProxy):
 
+    def __init__(self, wrapped, generator=None):
+        super().__init__(wrapped)
+        self._self_generator = generator
+
     @property
     def co_flags(self):
-        return self.__wrapped__.co_flags & _CO_SYNC_MASK
+        original = self.__wrapped__.co_flags
+        # Strip async-axis and iterable-coroutine bits; sync means neither
+        # coroutine function nor async generator nor types.coroutine-style.
+        flags = original & ~(CO_COROUTINE | CO_ASYNC_GENERATOR | CO_ITERABLE_COROUTINE)
+        if self._self_generator is True:
+            flags |= CO_GENERATOR
+        elif self._self_generator is False:
+            flags &= ~CO_GENERATOR
+        else:
+            # Auto: if input was an async generator, preserve generator-ness
+            # on the sync side by setting CO_GENERATOR. Otherwise leave
+            # CO_GENERATOR as-is (already copied from the wrapped flags).
+            if original & CO_ASYNC_GENERATOR:
+                flags |= CO_GENERATOR
+        return flags
 
 
 class _SyncFunctionSurrogate(CallableObjectProxy):
 
+    def __init__(self, wrapped, generator=None):
+        super().__init__(wrapped)
+        self._self_generator = generator
+
     @property
     def __code__(self):
-        return _SyncCodeProxy(self.__wrapped__.__code__)
+        return _SyncCodeProxy(self.__wrapped__.__code__, self._self_generator)
 
 
 class _BoundSyncFunctionWrapper(BoundFunctionWrapper):
@@ -48,77 +73,151 @@ def __init__(self, *args, **kwargs):
 
     @property
     def __func__(self):
-        return _SyncFunctionSurrogate(self.__wrapped__.__func__)
+        return _SyncFunctionSurrogate(
+            self.__wrapped__.__func__, self._self_parent._self_generator
+        )
 
 
 class _SyncFunctionWrapper(FunctionWrapper):
 
     __bound_function_wrapper__ = _BoundSyncFunctionWrapper
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, wrapped, wrapper, generator=None):
+        super().__init__(wrapped, wrapper)
         self._self_is_not_coroutine = True
+        self._self_generator = generator
 
     @property
     def __code__(self):
-        return _SyncCodeProxy(self.__wrapped__.__code__)
+        return _SyncCodeProxy(self.__wrapped__.__code__, self._self_generator)
 
 
 class _AsyncCodeProxy(CallableObjectProxy):
 
+    def __init__(self, wrapped, generator=None):
+        super().__init__(wrapped)
+        self._self_generator = generator
+
     @property
     def co_flags(self):
-        return (self.__wrapped__.co_flags & ~CO_ASYNC_GENERATOR) | CO_COROUTINE
+        original = self.__wrapped__.co_flags
+        # Strip all four convention bits; we reassert the right ones below.
+        flags = original & ~(
+            CO_GENERATOR | CO_COROUTINE | CO_ITERABLE_COROUTINE | CO_ASYNC_GENERATOR
+        )
+        if self._self_generator is True:
+            flags |= CO_ASYNC_GENERATOR
+        elif self._self_generator is False:
+            flags |= CO_COROUTINE
+        else:
+            # Auto: if input was a generator (sync or async), produce an
+            # async generator; otherwise produce a coroutine function.
+            if original & (CO_GENERATOR | CO_ASYNC_GENERATOR):
+                flags |= CO_ASYNC_GENERATOR
+            else:
+                flags |= CO_COROUTINE
+        return flags
 
 
 class _AsyncFunctionSurrogate(CallableObjectProxy):
 
+    def __init__(self, wrapped, generator=None):
+        super().__init__(wrapped)
+        self._self_generator = generator
+
     @property
     def __code__(self):
-        return _AsyncCodeProxy(self.__wrapped__.__code__)
+        return _AsyncCodeProxy(self.__wrapped__.__code__, self._self_generator)
 
 
 class _BoundAsyncFunctionWrapper(BoundFunctionWrapper):
 
     @property
     def __func__(self):
-        return _AsyncFunctionSurrogate(self.__wrapped__.__func__)
+        return _AsyncFunctionSurrogate(
+            self.__wrapped__.__func__, self._self_parent._self_generator
+        )
 
 
 class _AsyncFunctionWrapper(FunctionWrapper):
 
     __bound_function_wrapper__ = _BoundAsyncFunctionWrapper
 
+    def __init__(self, wrapped, wrapper, generator=None):
+        super().__init__(wrapped, wrapper)
+        self._self_generator = generator
+
     @property
     def __code__(self):
-        return _AsyncCodeProxy(self.__wrapped__.__code__)
+        return _AsyncCodeProxy(self.__wrapped__.__code__, self._self_generator)
 
 
-def mark_as_sync(wrapped):
+def mark_as_sync(wrapped=None, /, *, generator=None):
     """Mark a callable as synchronous from the perspective of calling
     convention detection. The returned wrapper is a pass-through that
     reports `inspect.iscoroutinefunction()` as False regardless of
     whether the underlying callable is declared `async def`. Useful
     when a stacked decorator has already collapsed an async function
-    into a synchronous one (for example by using `asyncio.run()`)."""
+    into a synchronous one (for example by using `asyncio.run()`).
 
-    def wrapper(wrapped, instance, args, kwargs):
-        return wrapped(*args, **kwargs)
+    The `generator` keyword toggles the sync generator bit
+    (`CO_GENERATOR`) on the resulting wrapper. Tri-state:
 
-    return _SyncFunctionWrapper(wrapped, wrapper)
+    - `None` (default): auto. Preserve generator-ness from the input --
+      if the input was an async generator, the wrapper reports as a sync
+      generator; otherwise CO_GENERATOR is copied through unchanged.
+    - `True`: force CO_GENERATOR on. Wrapper reports as a sync generator.
+    - `False`: force CO_GENERATOR off. Wrapper reports as a plain sync
+      function even if the input had CO_GENERATOR set.
+
+    Regardless of `generator`, CO_COROUTINE, CO_ASYNC_GENERATOR, and
+    CO_ITERABLE_COROUTINE are all cleared (sync means none of those).
+    """
+
+    def _decorator(wrapped):
+        def _wrapper(wrapped, instance, args, kwargs):
+            return wrapped(*args, **kwargs)
 
+        return _SyncFunctionWrapper(wrapped, _wrapper, generator=generator)
 
-def mark_as_async(wrapped):
+    if wrapped is None:
+        return _decorator
+    return _decorator(wrapped)
+
+
+def mark_as_async(wrapped=None, /, *, generator=None):
     """Mark a callable as asynchronous from the perspective of calling
     convention detection. The returned wrapper reports
     `inspect.iscoroutinefunction()` as True regardless of whether the
     underlying callable is declared `async def`. Useful when a stacked
-    decorator returns a coroutine from a plain `def` wrapper."""
+    decorator returns a coroutine from a plain `def` wrapper.
+
+    The `generator` keyword chooses between coroutine function and
+    async generator reporting. Tri-state:
+
+    - `None` (default): auto. If the input was a sync or async
+      generator, the wrapper reports as an async generator
+      (`CO_ASYNC_GENERATOR`); otherwise it reports as a coroutine
+      function (`CO_COROUTINE`).
+    - `True`: force async generator reporting (`CO_ASYNC_GENERATOR` set,
+      `CO_COROUTINE` cleared). These two flags are mutually exclusive at
+      the CPython code-object level.
+    - `False`: force coroutine function reporting (`CO_COROUTINE` set,
+      `CO_ASYNC_GENERATOR` cleared).
+
+    CO_GENERATOR and CO_ITERABLE_COROUTINE are always cleared (the
+    async path does not use either).
+    """
 
-    async def wrapper(wrapped, instance, args, kwargs):
+    async def _wrapper(wrapped, instance, args, kwargs):
         return wrapped(*args, **kwargs)
 
-    return _AsyncFunctionWrapper(wrapped, wrapper)
+    def _decorator(wrapped):
+        return _AsyncFunctionWrapper(wrapped, _wrapper, generator=generator)
+
+    if wrapped is None:
+        return _decorator
+    return _decorator(wrapped)
 
 
 def async_to_sync(wrapped):