Improve ext.asyncio.engine (#51)

* Improve `ext.asyncio.engine`

* Updates based on PR feedback

Author: bryanforbes

sqlalchemy-stubs/ext/asyncio/engine.pyi

@@ -1,103 +1,128 @@
 from typing import Any
 from typing import Callable
+from typing import Generator
 from typing import Mapping
+from typing import NoReturn
 from typing import Optional
+from typing import TypeVar
 
-from .base import ProxyComparable as ProxyComparable
-from .base import StartableContext as StartableContext
-from .result import AsyncResult as AsyncResult
-from ... import exc as exc
-from ... import util as util
-from ...engine import Result as Result
-from ...engine import Transaction as Transaction
-from ...future import Connection as Connection
-from ...future import Engine as Engine
-from ...sql import Executable as Executable
-from ...util.concurrency import greenlet_spawn as greenlet_spawn
+from .base import ProxyComparable
+from .base import StartableContext
+from .result import AsyncResult
+from ...engine import Result
+from ...engine import Transaction
+from ...future import Connection
+from ...future import Engine
+from ...sql import Executable
 
-def create_async_engine(*arg: Any, **kw: Any): ...
+_TAsyncConnection = TypeVar("_TAsyncConnection", bound=AsyncConnection)
+_TAsyncTransaction = TypeVar("_TAsyncTransaction", bound=AsyncTransaction)
+
+def create_async_engine(*arg: Any, **kw: Any) -> AsyncEngine: ...
 
 class AsyncConnectable: ...
 
 class AsyncConnection(ProxyComparable, StartableContext, AsyncConnectable):
-    engine: Any = ...
-    sync_engine: Any = ...
-    sync_connection: Any = ...
+    # copied from future.Connection via create_proxy_methods
+    @property
+    def closed(self) -> bool: ...
+    @property
+    def invalidated(self) -> bool: ...
+    dialect: Any
+    @property
+    def default_isolation_level(self) -> Any: ...
+    # end copied
+
+    engine: AsyncEngine = ...
+    sync_engine: Engine = ...
+    sync_connection: Optional[Connection] = ...
     def __init__(
         self,
         async_engine: AsyncEngine,
         sync_connection: Optional[Connection] = ...,
     ) -> None: ...
-    async def start(self): ...
+    async def start(self: _TAsyncConnection) -> _TAsyncConnection: ...
     @property
-    def connection(self) -> None: ...
-    async def get_raw_connection(self): ...
+    def connection(self) -> NoReturn: ...
+    async def get_raw_connection(self) -> Any: ...
     @property
-    def info(self): ...
+    def info(self) -> Mapping[Any, Any]: ...
     def begin(self) -> AsyncTransaction: ...
     def begin_nested(self) -> AsyncTransaction: ...
-    async def invalidate(self, exception: Optional[Any] = ...): ...
-    async def get_isolation_level(self): ...
-    async def set_isolation_level(self): ...
-    def in_transaction(self): ...
-    def in_nested_transaction(self): ...
-    def get_transaction(self): ...
-    def get_nested_transaction(self): ...
-    async def execution_options(self, **opt: Any): ...
+    async def invalidate(self, exception: Optional[Any] = ...) -> None: ...
+    async def get_isolation_level(self) -> Any: ...
+    async def set_isolation_level(self) -> Any: ...
+    def in_transaction(self) -> bool: ...
+    def in_nested_transaction(self) -> bool: ...
+    def get_transaction(self) -> Optional[AsyncTransaction]: ...
+    def get_nested_transaction(self) -> Optional[AsyncTransaction]: ...
+    async def execution_options(
+        self: _TAsyncConnection, **opt: Any
+    ) -> _TAsyncConnection: ...
     async def commit(self) -> None: ...
     async def rollback(self) -> None: ...
     async def close(self) -> None: ...
     async def exec_driver_sql(
         self,
-        statement: Executable,
-        parameters: Optional[Mapping] = ...,
-        execution_options: Mapping = ...,
+        statement: str,
+        parameters: Optional[Mapping[Any, Any]] = ...,
+        execution_options: Mapping[Any, Any] = ...,
     ) -> Result: ...
     async def stream(
         self,
         statement: Executable,
-        parameters: Optional[Mapping] = ...,
-        execution_options: Mapping = ...,
+        parameters: Optional[Mapping[Any, Any]] = ...,
+        execution_options: Mapping[Any, Any] = ...,
     ) -> AsyncResult: ...
     async def execute(
         self,
         statement: Executable,
-        parameters: Optional[Mapping] = ...,
-        execution_options: Mapping = ...,
+        parameters: Optional[Mapping[Any, Any]] = ...,
+        execution_options: Mapping[Any, Any] = ...,
     ) -> Result: ...
     async def scalar(
         self,
         statement: Executable,
-        parameters: Optional[Mapping] = ...,
-        execution_options: Mapping = ...,
+        parameters: Optional[Mapping[Any, Any]] = ...,
+        execution_options: Mapping[Any, Any] = ...,
+    ) -> Any: ...
+    async def run_sync(
+        self, fn: Callable[..., Any], *arg: Any, **kw: Any
     ) -> Any: ...
-    async def run_sync(self, fn: Callable, *arg: Any, **kw: Any) -> Any: ...
-    def __await__(self): ...
+    def __await__(
+        self: _TAsyncConnection,
+    ) -> Generator[Any, None, _TAsyncConnection]: ...
     async def __aexit__(
         self, type_: Any, value: Any, traceback: Any
     ) -> None: ...
 
 class AsyncEngine(ProxyComparable, AsyncConnectable):
+    # copied from future.Engine by create_proxy_methods
+    def clear_compiled_cache(self) -> None: ...
+    def update_execution_options(self, **opt: Any) -> None: ...
+    def get_execution_options(self) -> Mapping[Any, Any]: ...
+    # end copied
     class _trans_ctx(StartableContext):
-        conn: Any = ...
-        def __init__(self, conn: Any) -> None: ...
+        conn: AsyncConnection = ...
+        def __init__(self, conn: AsyncConnection) -> None: ...
         transaction: Any = ...
-        async def start(self): ...
+        async def start(self) -> AsyncConnection: ...  # type: ignore[override]
+        def __await__(self) -> Generator[Any, None, AsyncConnection]: ...  # type: ignore[override]
         async def __aexit__(
             self, type_: Any, value: Any, traceback: Any
         ) -> None: ...
-    sync_engine: Any = ...
+    sync_engine: Engine = ...
     def __init__(self, sync_engine: Engine) -> None: ...
-    def begin(self): ...
+    def begin(self) -> _trans_ctx: ...
     def connect(self) -> AsyncConnection: ...
     async def raw_connection(self) -> Any: ...
-    def execution_options(self, **opt: Any): ...
-    async def dispose(self): ...
+    def execution_options(self, **opt: Any) -> AsyncEngine: ...
+    async def dispose(self) -> None: ...
 
 class AsyncTransaction(ProxyComparable, StartableContext):
-    connection: Any = ...
-    sync_transaction: Any = ...
-    nested: Any = ...
+    connection: AsyncConnection = ...
+    sync_transaction: Optional[Transaction] = ...
+    nested: bool = ...
     def __init__(
         self, connection: AsyncConnection, nested: bool = ...
     ) -> None: ...
@@ -108,7 +133,7 @@ class AsyncTransaction(ProxyComparable, StartableContext):
     async def close(self) -> None: ...
     async def rollback(self) -> None: ...
     async def commit(self) -> None: ...
-    async def start(self): ...
+    async def start(self: _TAsyncTransaction) -> _TAsyncTransaction: ...
     async def __aexit__(
         self, type_: Any, value: Any, traceback: Any
     ) -> None: ...