Проверка 09.02.2025

venv/Lib/site-packages/sqlalchemy/util/_concurrency_py3k.py

@@ -1,5 +1,5 @@
 # util/_concurrency_py3k.py
-# Copyright (C) 2005-2024 the SQLAlchemy authors and contributors
+# Copyright (C) 2005-2025 the SQLAlchemy authors and contributors
 # <see AUTHORS file>
 #
 # This module is part of SQLAlchemy and is released under
@@ -19,10 +19,14 @@ from typing import Coroutine
 from typing import Optional
 from typing import TYPE_CHECKING
 from typing import TypeVar
+from typing import Union
 
 from .langhelpers import memoized_property
 from .. import exc
+from ..util import py311
+from ..util.typing import Literal
 from ..util.typing import Protocol
+from ..util.typing import Self
 from ..util.typing import TypeGuard
 
 _T = TypeVar("_T")
@@ -70,9 +74,10 @@ def is_exit_exception(e: BaseException) -> bool:
 class _AsyncIoGreenlet(greenlet):
     dead: bool
 
+    __sqlalchemy_greenlet_provider__ = True
+
     def __init__(self, fn: Callable[..., Any], driver: greenlet):
         greenlet.__init__(self, fn, driver)
-        self.driver = driver
         if _has_gr_context:
             self.gr_context = driver.gr_context
 
@@ -98,7 +103,7 @@ def _safe_cancel_awaitable(awaitable: Awaitable[Any]) -> None:
 
 def in_greenlet() -> bool:
     current = getcurrent()
-    return isinstance(current, _AsyncIoGreenlet)
+    return getattr(current, "__sqlalchemy_greenlet_provider__", False)
 
 
 def await_only(awaitable: Awaitable[_T]) -> _T:
@@ -112,7 +117,7 @@ def await_only(awaitable: Awaitable[_T]) -> _T:
     """
     # this is called in the context greenlet while running fn
     current = getcurrent()
-    if not isinstance(current, _AsyncIoGreenlet):
+    if not getattr(current, "__sqlalchemy_greenlet_provider__", False):
         _safe_cancel_awaitable(awaitable)
 
         raise exc.MissingGreenlet(
@@ -124,7 +129,7 @@ def await_only(awaitable: Awaitable[_T]) -> _T:
     # a coroutine to run. Once the awaitable is done, the driver greenlet
     # switches back to this greenlet with the result of awaitable that is
     # then returned to the caller (or raised as error)
-    return current.driver.switch(awaitable)  # type: ignore[no-any-return]
+    return current.parent.switch(awaitable)  # type: ignore[no-any-return,attr-defined] # noqa: E501
 
 
 def await_fallback(awaitable: Awaitable[_T]) -> _T:
@@ -144,7 +149,7 @@ def await_fallback(awaitable: Awaitable[_T]) -> _T:
 
     # this is called in the context greenlet while running fn
     current = getcurrent()
-    if not isinstance(current, _AsyncIoGreenlet):
+    if not getattr(current, "__sqlalchemy_greenlet_provider__", False):
         loop = get_event_loop()
         if loop.is_running():
             _safe_cancel_awaitable(awaitable)
@@ -156,7 +161,7 @@ def await_fallback(awaitable: Awaitable[_T]) -> _T:
             )
         return loop.run_until_complete(awaitable)
 
-    return current.driver.switch(awaitable)  # type: ignore[no-any-return]
+    return current.parent.switch(awaitable)  # type: ignore[no-any-return,attr-defined] # noqa: E501
 
 
 async def greenlet_spawn(
@@ -182,24 +187,21 @@ async def greenlet_spawn(
     # coroutine to wait. If the context is dead the function has
     # returned, and its result can be returned.
     switch_occurred = False
-    try:
-        result = context.switch(*args, **kwargs)
-        while not context.dead:
-            switch_occurred = True
-            try:
-                # wait for a coroutine from await_only and then return its
-                # result back to it.
-                value = await result
-            except BaseException:
-                # this allows an exception to be raised within
-                # the moderated greenlet so that it can continue
-                # its expected flow.
-                result = context.throw(*sys.exc_info())
-            else:
-                result = context.switch(value)
-    finally:
-        # clean up to avoid cycle resolution by gc
-        del context.driver
+    result = context.switch(*args, **kwargs)
+    while not context.dead:
+        switch_occurred = True
+        try:
+            # wait for a coroutine from await_only and then return its
+            # result back to it.
+            value = await result
+        except BaseException:
+            # this allows an exception to be raised within
+            # the moderated greenlet so that it can continue
+            # its expected flow.
+            result = context.throw(*sys.exc_info())
+        else:
+            result = context.switch(value)
+
     if _require_await and not switch_occurred:
         raise exc.AwaitRequired(
             "The current operation required an async execution but none was "
@@ -225,34 +227,6 @@ class AsyncAdaptedLock:
         self.mutex.release()
 
 
-def _util_async_run_coroutine_function(
-    fn: Callable[..., Coroutine[Any, Any, Any]], *args: Any, **kwargs: Any
-) -> Any:
-    """for test suite/ util only"""
-
-    loop = get_event_loop()
-    if loop.is_running():
-        raise Exception(
-            "for async run coroutine we expect that no greenlet or event "
-            "loop is running when we start out"
-        )
-    return loop.run_until_complete(fn(*args, **kwargs))
-
-
-def _util_async_run(
-    fn: Callable[..., Coroutine[Any, Any, Any]], *args: Any, **kwargs: Any
-) -> Any:
-    """for test suite/ util only"""
-
-    loop = get_event_loop()
-    if not loop.is_running():
-        return loop.run_until_complete(greenlet_spawn(fn, *args, **kwargs))
-    else:
-        # allow for a wrapped test function to call another
-        assert isinstance(getcurrent(), _AsyncIoGreenlet)
-        return fn(*args, **kwargs)
-
-
 def get_event_loop() -> asyncio.AbstractEventLoop:
     """vendor asyncio.get_event_loop() for python 3.7 and above.
 
@@ -265,3 +239,50 @@ def get_event_loop() -> asyncio.AbstractEventLoop:
         # avoid "During handling of the above exception, another exception..."
         pass
     return asyncio.get_event_loop_policy().get_event_loop()
+
+
+if not TYPE_CHECKING and py311:
+    _Runner = asyncio.Runner
+else:
+
+    class _Runner:
+        """Runner implementation for test only"""
+
+        _loop: Union[None, asyncio.AbstractEventLoop, Literal[False]]
+
+        def __init__(self) -> None:
+            self._loop = None
+
+        def __enter__(self) -> Self:
+            self._lazy_init()
+            return self
+
+        def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+            self.close()
+
+        def close(self) -> None:
+            if self._loop:
+                try:
+                    self._loop.run_until_complete(
+                        self._loop.shutdown_asyncgens()
+                    )
+                finally:
+                    self._loop.close()
+                    self._loop = False
+
+        def get_loop(self) -> asyncio.AbstractEventLoop:
+            """Return embedded event loop."""
+            self._lazy_init()
+            assert self._loop
+            return self._loop
+
+        def run(self, coro: Coroutine[Any, Any, _T]) -> _T:
+            self._lazy_init()
+            assert self._loop
+            return self._loop.run_until_complete(coro)
+
+        def _lazy_init(self) -> None:
+            if self._loop is False:
+                raise RuntimeError("Runner is closed")
+            if self._loop is None:
+                self._loop = asyncio.new_event_loop()