# ext/asyncio/base.py
# Copyright (C) 2020-2025 the SQLAlchemy authors and contributors
# <see AUTHORS file>
#
# This module is part of SQLAlchemy and is released under
# the MIT License: https://www.opensource.org/licenses/mit-license.php
from __future__ import annotations
import abc
import functools
from typing import Any
from typing import AsyncGenerator
from typing import AsyncIterator
from typing import Awaitable
from typing import Callable
from typing import ClassVar
from typing import Dict
from typing import Generator
from typing import Generic
from typing import NoReturn
from typing import Optional
from typing import overload
from typing import Tuple
from typing import TypeVar
import weakref
from . import exc as async_exc
from ... import util
from ...util.typing import Literal
from ...util.typing import Self
_T = TypeVar("_T", bound=Any)
_T_co = TypeVar("_T_co", bound=Any, covariant=True)
_PT = TypeVar("_PT", bound=Any)
class ReversibleProxy(Generic[_PT]):
_proxy_objects: ClassVar[
Dict[weakref.ref[Any], weakref.ref[ReversibleProxy[Any]]]
] = {}
__slots__ = ("__weakref__",)
@overload
def _assign_proxied(self, target: _PT) -> _PT: ...
@overload
def _assign_proxied(self, target: None) -> None: ...
def _assign_proxied(self, target: Optional[_PT]) -> Optional[_PT]:
if target is not None:
target_ref: weakref.ref[_PT] = weakref.ref(
target, ReversibleProxy._target_gced
)
proxy_ref = weakref.ref(
self,
functools.partial(ReversibleProxy._target_gced, target_ref),
)
ReversibleProxy._proxy_objects[target_ref] = proxy_ref
return target
@classmethod
def _target_gced(
cls,
ref: weakref.ref[_PT],
proxy_ref: Optional[weakref.ref[Self]] = None, # noqa: U100
) -> None:
cls._proxy_objects.pop(ref, None)
@classmethod
def _regenerate_proxy_for_target(cls, target: _PT) -> Self:
raise NotImplementedError()
@overload
@classmethod
def _retrieve_proxy_for_target(
cls,
target: _PT,
regenerate: Literal[True] = ...,
) -> Self: ...
@overload
@classmethod
def _retrieve_proxy_for_target(
cls, target: _PT, regenerate: bool = True
) -> Optional[Self]: ...
@classmethod
def _retrieve_proxy_for_target(
cls, target: _PT, regenerate: bool = True
) -> Optional[Self]:
try:
proxy_ref = cls._proxy_objects[weakref.ref(target)]
except KeyError:
pass
else:
proxy = proxy_ref()
if proxy is not None:
return proxy # type: ignore
if regenerate:
return cls._regenerate_proxy_for_target(target)
else:
return None
class StartableContext(Awaitable[_T_co], abc.ABC):
__slots__ = ()
@abc.abstractmethod
async def start(self, is_ctxmanager: bool = False) -> _T_co:
raise NotImplementedError()
def __await__(self) -> Generator[Any, Any, _T_co]:
return self.start().__await__()
async def __aenter__(self) -> _T_co:
return await self.start(is_ctxmanager=True)
@abc.abstractmethod
async def __aexit__(
self, type_: Any, value: Any, traceback: Any
) -> Optional[bool]:
pass
def _raise_for_not_started(self) -> NoReturn:
raise async_exc.AsyncContextNotStarted(
"%s context has not been started and object has not been awaited."
% (self.__class__.__name__)
)
class GeneratorStartableContext(StartableContext[_T_co]):
__slots__ = ("gen",)
gen: AsyncGenerator[_T_co, Any]
def __init__(
self,
func: Callable[..., AsyncIterator[_T_co]],
args: Tuple[Any, ...],
kwds: Dict[str, Any],
):
self.gen = func(*args, **kwds) # type: ignore
async def start(self, is_ctxmanager: bool = False) -> _T_co:
try:
start_value = await util.anext_(self.gen)
except StopAsyncIteration:
raise RuntimeError("generator didn't yield") from None
# if not a context manager, then interrupt the generator, don't
# let it complete. this step is technically not needed, as the
# generator will close in any case at gc time. not clear if having
# this here is a good idea or not (though it helps for clarity IMO)
if not is_ctxmanager:
await self.gen.aclose()
return start_value
async def __aexit__(
self, typ: Any, value: Any, traceback: Any
) -> Optional[bool]:
# vendored from contextlib.py
if typ is None:
try:
await util.anext_(self.gen)
except StopAsyncIteration:
return False
else:
raise RuntimeError("generator didn't stop")
else:
if value is None:
# Need to force instantiation so we can reliably
# tell if we get the same exception back
value = typ()
try:
await self.gen.athrow(value)
except StopAsyncIteration as exc:
# Suppress StopIteration *unless* it's the same exception that
# was passed to throw(). This prevents a StopIteration
# raised inside the "with" statement from being suppressed.
return exc is not value
except RuntimeError as exc:
# Don't re-raise the passed in exception. (issue27122)
if exc is value:
return False
# Avoid suppressing if a Stop(Async)Iteration exception
# was passed to athrow() and later wrapped into a RuntimeError
# (see PEP 479 for sync generators; async generators also
# have this behavior). But do this only if the exception
# wrapped
# by the RuntimeError is actully Stop(Async)Iteration (see
# issue29692).
if (
isinstance(value, (StopIteration, StopAsyncIteration))
and exc.__cause__ is value
):
return False
raise
except BaseException as exc:
# only re-raise if it's *not* the exception that was
# passed to throw(), because __exit__() must not raise
# an exception unless __exit__() itself failed. But throw()
# has to raise the exception to signal propagation, so this
# fixes the impedance mismatch between the throw() protocol
# and the __exit__() protocol.
if exc is not value:
raise
return False
raise RuntimeError("generator didn't stop after athrow()")
def asyncstartablecontext(
func: Callable[..., AsyncIterator[_T_co]]
) -> Callable[..., GeneratorStartableContext[_T_co]]:
"""@asyncstartablecontext decorator.
the decorated function can be called either as ``async with fn()``, **or**
``await fn()``. This is decidedly different from what
``@contextlib.asynccontextmanager`` supports, and the usage pattern
is different as well.
Typical usage:
.. sourcecode:: text
@asyncstartablecontext
async def some_async_generator(<arguments>):
<setup>
try:
yield <value>
except GeneratorExit:
# return value was awaited, no context manager is present
# and caller will .close() the resource explicitly
pass
else:
<context manager cleanup>
Above, ``GeneratorExit`` is caught if the function were used as an
``await``. In this case, it's essential that the cleanup does **not**
occur, so there should not be a ``finally`` block.
If ``GeneratorExit`` is not invoked, this means we're in ``__aexit__``
and we were invoked as a context manager, and cleanup should proceed.
"""
@functools.wraps(func)
def helper(*args: Any, **kwds: Any) -> GeneratorStartableContext[_T_co]:
return GeneratorStartableContext(func, args, kwds)
return helper
class ProxyComparable(ReversibleProxy[_PT]):
__slots__ = ()
@util.ro_non_memoized_property
def _proxied(self) -> _PT:
raise NotImplementedError()
def __hash__(self) -> int:
return id(self)
def __eq__(self, other: Any) -> bool:
return (
isinstance(other, self.__class__)
and self._proxied == other._proxied
)
def __ne__(self, other: Any) -> bool:
return (
not isinstance(other, self.__class__)
or self._proxied != other._proxied
)
|