Merge "add contextmanager typing, open run_sync typing" into main

2 files changed, 9 insertions, 14 deletions

diff --git a/lib/sqlalchemy/ext/asyncio/engine.py b/lib/sqlalchemy/ext/asyncio/engine.py
index 97d69fcbd..2418dab88 100644
--- a/lib/sqlalchemy/ext/asyncio/engine.py
+++ b/lib/sqlalchemy/ext/asyncio/engine.py
@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import Generator
 from typing import NoReturn
@@ -33,7 +34,6 @@ from ...engine import Engine
 from ...engine.base import NestedTransaction
 from ...engine.base import Transaction
 from ...util.concurrency import greenlet_spawn
-from ...util.typing import Protocol
 
 if TYPE_CHECKING:
     from ...engine.cursor import CursorResult
@@ -55,11 +55,6 @@ if TYPE_CHECKING:
 _T = TypeVar("_T", bound=Any)
 
 
-class _SyncConnectionCallable(Protocol):
-    def __call__(self, connection: Connection, *arg: Any, **kw: Any) -> Any:
-        ...
-
-
 def create_async_engine(url: Union[str, URL], **kw: Any) -> AsyncEngine:
     """Create a new async engine instance.
 
@@ -667,7 +662,7 @@ class AsyncConnection(
         return result.scalars()
 
     async def run_sync(
-        self, fn: _SyncConnectionCallable, *arg: Any, **kw: Any
+        self, fn: Callable[..., Any], *arg: Any, **kw: Any
     ) -> Any:
         """Invoke the given sync callable passing self as the first argument.
 
@@ -845,6 +840,11 @@ class AsyncEngine(ProxyComparable[Engine], AsyncConnectable):
         def __init__(self, conn: AsyncConnection):
             self.conn = conn
 
+        if TYPE_CHECKING:
+
+            async def __aenter__(self) -> AsyncConnection:
+                ...
+
         async def start(self, is_ctxmanager: bool = False) -> AsyncConnection:
             await self.conn.start(is_ctxmanager=is_ctxmanager)
             self.transaction = self.conn.begin()
diff --git a/lib/sqlalchemy/ext/asyncio/session.py b/lib/sqlalchemy/ext/asyncio/session.py
index be3414cef..ea587f8de 100644
--- a/lib/sqlalchemy/ext/asyncio/session.py
+++ b/lib/sqlalchemy/ext/asyncio/session.py
@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import Generic
 from typing import Iterable
@@ -33,7 +34,6 @@ from ...orm import Session
 from ...orm import SessionTransaction
 from ...orm import state as _instance_state
 from ...util.concurrency import greenlet_spawn
-from ...util.typing import Protocol
 
 if TYPE_CHECKING:
     from .engine import AsyncConnection
@@ -68,11 +68,6 @@ _AsyncSessionBind = Union["AsyncEngine", "AsyncConnection"]
 _T = TypeVar("_T", bound=Any)
 
 
-class _SyncSessionCallable(Protocol):
-    def __call__(self, session: Session, *arg: Any, **kw: Any) -> Any:
-        ...
-
-
 _EXECUTE_OPTIONS = util.immutabledict({"prebuffer_rows": True})
 _STREAM_OPTIONS = util.immutabledict({"stream_results": True})
 
@@ -234,7 +229,7 @@ class AsyncSession(ReversibleProxy[Session]):
         )
 
     async def run_sync(
-        self, fn: _SyncSessionCallable, *arg: Any, **kw: Any
+        self, fn: Callable[..., Any], *arg: Any, **kw: Any
     ) -> Any:
         """Invoke the given sync callable passing sync self as the first
         argument.