Improve ext.asyncio.session (#52)

* Improve `ext.asyncio.session`

* Add `AsyncSession.delete()` to stubs

* Fix bind and get_bind types

Author: bryanforbes

sqlalchemy-stubs/ext/asyncio/session.pyi

@@ -1,54 +1,72 @@
 from typing import Any
 from typing import Callable
+from typing import ContextManager
+from typing import Iterable
+from typing import Iterator
 from typing import Mapping
 from typing import Optional
+from typing import Sequence
 from typing import TypeVar
+from typing import Union
 
-from . import engine as engine
-from .base import StartableContext as StartableContext
-from .engine import AsyncEngine as AsyncEngine
-from ... import util as util
-from ...engine import Result as Result
-from ...orm import Session as Session
-from ...sql import Executable as Executable
-from ...util.concurrency import greenlet_spawn as greenlet_spawn
+from .base import StartableContext
+from .engine import AsyncConnection
+from .engine import AsyncEngine
+from .result import AsyncResult
+from ... import util
+from ...engine import Result
+from ...engine.base import _ExecutionOptions
+from ...orm import Session
+from ...orm.session import _IdentityMap
+from ...sql import ClauseElement
+from ...sql import Executable
 
 _T = TypeVar("_T")
+_TAsyncSession = TypeVar("_TAsyncSession", bound=AsyncSession)
+_TAsyncSessionTransaction = TypeVar(
+    "_TAsyncSessionTransaction", bound=AsyncSessionTransaction
+)
 
 class AsyncSession:
     dispatch: Any = ...
     bind: Any = ...
     binds: Any = ...
-    sync_session: Any = ...
+    sync_session: Session = ...
     def __init__(
         self,
-        bind: AsyncEngine = ...,
-        binds: Mapping[object, AsyncEngine] = ...,
+        bind: Optional[Union[AsyncConnection, AsyncEngine]] = ...,
+        binds: Optional[
+            Mapping[object, Union[AsyncConnection, AsyncEngine]]
+        ] = ...,
         **kw: Any,
     ) -> None: ...
     async def refresh(
         self,
         instance: Any,
         attribute_names: Optional[Any] = ...,
         with_for_update: Optional[Any] = ...,
-    ): ...
+    ) -> None: ...
     async def run_sync(
         self, fn: Callable[..., _T], *arg: Any, **kw: Any
     ) -> _T: ...
     async def execute(
         self,
         statement: Executable,
-        params: Optional[Mapping] = ...,
-        execution_options: Mapping = ...,
-        bind_arguments: Optional[Mapping] = ...,
+        params: Optional[
+            Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]
+        ] = ...,
+        execution_options: Optional[_ExecutionOptions] = ...,
+        bind_arguments: Optional[Mapping[str, Any]] = ...,
         **kw: Any,
     ) -> Result: ...
     async def scalar(
         self,
         statement: Executable,
-        params: Optional[Mapping] = ...,
-        execution_options: Mapping = ...,
-        bind_arguments: Optional[Mapping] = ...,
+        params: Optional[
+            Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]
+        ] = ...,
+        execution_options: _ExecutionOptions = ...,
+        bind_arguments: Optional[Mapping[str, Any]] = ...,
         **kw: Any,
     ) -> Any: ...
     async def get(
@@ -59,47 +77,93 @@ class AsyncSession:
         populate_existing: bool = ...,
         with_for_update: Optional[Any] = ...,
         identity_token: Optional[Any] = ...,
-    ): ...
+    ) -> Any: ...
     async def stream(
         self,
         statement: Any,
-        params: Optional[Any] = ...,
-        execution_options: Any = ...,
-        bind_arguments: Optional[Any] = ...,
+        params: Optional[
+            Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]
+        ] = ...,
+        execution_options: _ExecutionOptions = ...,
+        bind_arguments: Optional[Mapping[str, Any]] = ...,
         **kw: Any,
-    ): ...
-    async def merge(self, instance: Any, load: bool = ...): ...
+    ) -> AsyncResult: ...
+    async def delete(self, instance: Any) -> None: ...
+    async def merge(self, instance: _T, load: bool = ...) -> _T: ...
     async def flush(self, objects: Optional[Any] = ...) -> None: ...
-    async def connection(self): ...
-    def begin(self, **kw: Any): ...
-    def begin_nested(self, **kw: Any): ...
-    async def rollback(self): ...
-    async def commit(self): ...
-    async def close(self): ...
+    async def connection(self) -> AsyncConnection: ...
+    def begin(self, **kw: Any) -> AsyncSessionTransaction: ...
+    def begin_nested(self, **kw: Any) -> AsyncSessionTransaction: ...
+    async def rollback(self) -> None: ...
+    async def commit(self) -> None: ...
+    async def close(self) -> None: ...
     @classmethod
-    async def close_all(self): ...
-    async def __aenter__(self): ...
+    async def close_all(self) -> None: ...
+    async def __aenter__(self: _TAsyncSession) -> _TAsyncSession: ...
     async def __aexit__(
         self, type_: Any, value: Any, traceback: Any
     ) -> None: ...
+    # copied via create_proxy_methods
+    def __contains__(self, instance: Any) -> bool: ...
+    def __iter__(self) -> Iterator[Any]: ...
+    def add(self, instance: Any, _warn: bool = ...) -> None: ...
+    def add_all(self, instances: Any) -> None: ...
+    def expire(
+        self, instance: Any, attribute_names: Optional[Iterable[str]] = ...
+    ) -> None: ...
+    def expire_all(self) -> None: ...
+    def expunge(self, instance: Any) -> None: ...
+    def expunge_all(self) -> None: ...
+    def get_bind(
+        self,
+        mapper: Optional[Any] = ...,
+        clause: Optional[ClauseElement] = ...,
+        bind: Optional[Union[AsyncConnection, AsyncEngine]] = ...,
+        _sa_skip_events: Optional[Any] = ...,
+        _sa_skip_for_implicit_returning: bool = ...,
+    ) -> Union[AsyncConnection, AsyncEngine]: ...
+    def is_modified(
+        self, instance: Any, include_collections: bool = ...
+    ) -> bool: ...
+    def in_transaction(self) -> bool: ...
+    @property
+    def dirty(self) -> util.IdentitySet[Any]: ...
+    @property
+    def deleted(self) -> util.IdentitySet[Any]: ...
+    @property
+    def new(self) -> util.IdentitySet[Any]: ...
+    identity_map: _IdentityMap
+    @property
+    def is_active(self) -> bool: ...
+    autoflush: bool
+    @property
+    def no_autoflush(
+        self: _TAsyncSession,
+    ) -> ContextManager[_TAsyncSession]: ...
+    @util.memoized_property
+    def info(self) -> Mapping[Any, Any]: ...
 
 class _AsyncSessionContextManager:
-    async_session: Any = ...
-    def __init__(self, async_session: Any) -> None: ...
-    trans: Any = ...
-    async def __aenter__(self): ...
+    async_session: AsyncSession = ...
+    trans: AsyncSessionTransaction = ...
+    def __init__(self, async_session: AsyncSession) -> None: ...
+    async def __aenter__(self) -> AsyncSession: ...
     async def __aexit__(
         self, type_: Any, value: Any, traceback: Any
     ) -> None: ...
 
 class AsyncSessionTransaction(StartableContext):
-    session: Any = ...
-    nested: Any = ...
-    sync_transaction: Any = ...
-    def __init__(self, session: Any, nested: bool = ...) -> None: ...
+    session: AsyncSession = ...
+    nested: bool = ...
+    sync_transaction: Optional[Any] = ...
+    def __init__(self, session: AsyncSession, nested: bool = ...) -> None: ...
     @property
-    def is_active(self): ...
-    async def rollback(self) -> None: ...
-    async def commit(self) -> None: ...
-    async def start(self): ...
-    async def __aexit__(self, type_: Any, value: Any, traceback: Any): ...
+    def is_active(self) -> bool: ...
+    async def rollback(self) -> Optional[AsyncSessionTransaction]: ...
+    async def commit(self) -> Optional[AsyncSessionTransaction]: ...
+    async def start(
+        self: _TAsyncSessionTransaction,
+    ) -> _TAsyncSessionTransaction: ...
+    async def __aexit__(
+        self, type_: Any, value: Any, traceback: Any
+    ) -> None: ...