Improve proxies. (#149)

* improve proxies.

* improve names of type only classes.

Author: CaselIT

setup.cfg

@@ -37,6 +37,7 @@ incremental = True
 strict = True
 warn_unused_ignores = False
 
+plugins = sqlalchemy.ext.mypy.plugin
 
 
 [flake8]

sqlalchemy-stubs/engine/base.pyi

@@ -36,9 +36,17 @@ class _ConnectionCallable(Protocol[_T_contra, _T_co]):
         self, __connection: _T_contra, *args: Any, **kwargs: Any
     ) -> _T_co: ...
 
-class Connection(Connectable):
-    engine: Engine = ...
+class _ConnectionTypingCommon:
+    @property
+    def closed(self) -> bool: ...
+    @property
+    def invalidated(self) -> bool: ...
     dialect: Dialect = ...
+    @property
+    def default_isolation_level(self) -> Any: ...
+
+class Connection(_ConnectionTypingCommon, Connectable):
+    engine: Engine = ...
     should_close_with_result: bool = ...
     dispatch: Any = ...
     def __init__(
@@ -57,15 +65,9 @@ class Connection(Connectable):
     def execution_options(self: _TConnection, **opt: Any) -> _TConnection: ...
     def get_execution_options(self) -> Dict[str, Any]: ...
     @property
-    def closed(self) -> bool: ...
-    @property
-    def invalidated(self) -> bool: ...
-    @property
     def connection(self) -> _DBAPIConnection: ...
     def get_isolation_level(self) -> Any: ...
     @property
-    def default_isolation_level(self) -> Any: ...
-    @property
     def info(self) -> MutableMapping[Any, Any]: ...
     def connect(self: _TConnection, close_with_result: bool = ...) -> _TConnection: ...  # type: ignore[override]
     def invalidate(self, exception: Optional[Any] = ...) -> None: ...
@@ -161,13 +163,25 @@ class TwoPhaseTransaction(RootTransaction):
     def __init__(self, connection: Connection, xid: str) -> None: ...
     def prepare(self) -> None: ...
 
-class Engine(Connectable, log.Identified):
+class _EngineTypingCommon:
     pool: Pool = ...
     url: URL = ...
     dialect: Dialect = ...
     logging_name: Optional[str] = ...
     echo: Optional[Union[bool, Literal["debug"]]] = ...
     hide_parameters: bool = ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def driver(self) -> str: ...
+    def clear_compiled_cache(self) -> None: ...
+    def update_execution_options(self, **opt: Any) -> None: ...
+    def get_execution_options(self) -> Dict[str, Any]: ...
+
+class Engine(_EngineTypingCommon, Connectable, log.Identified):
+    @property
+    def engine(self: _TEngine) -> _TEngine: ...
+    hide_parameters: bool = ...
     def __init__(
         self,
         pool: Pool,
@@ -179,16 +193,7 @@ class Engine(Connectable, log.Identified):
         execution_options: Optional[Dict[str, Any]] = ...,
         hide_parameters: bool = ...,
     ) -> None: ...
-    @property
-    def engine(self: _TEngine) -> _TEngine: ...  # type: ignore[override]
-    def clear_compiled_cache(self) -> None: ...
-    def update_execution_options(self, **opt: Any) -> None: ...
     def execution_options(self, **opt: Any) -> OptionEngine: ...
-    def get_execution_options(self) -> Dict[str, Any]: ...
-    @property
-    def name(self) -> str: ...
-    @property
-    def driver(self) -> str: ...
     def dispose(self) -> None: ...
     class _trans_ctx:
         conn: Connection = ...

sqlalchemy-stubs/ext/asyncio/__init__.pyi

@@ -7,5 +7,8 @@ from .events import AsyncSessionEvents as AsyncSessionEvents
 from .result import AsyncMappingResult as AsyncMappingResult
 from .result import AsyncResult as AsyncResult
 from .result import AsyncScalarResult as AsyncScalarResult
+from .scoping import async_scoped_session as async_scoped_session
+from .session import async_object_session as async_object_session
+from .session import async_session as async_session
 from .session import AsyncSession as AsyncSession
 from .session import AsyncSessionTransaction as AsyncSessionTransaction

sqlalchemy-stubs/ext/asyncio/engine.pyi

@@ -10,28 +10,35 @@ from typing import TypeVar
 from .base import ProxyComparable
 from .base import StartableContext
 from .result import AsyncResult
+from ...engine import Dialect
 from ...engine import Result
 from ...engine import Transaction
+from ...engine.base import _ConnectionTypingCommon
+from ...engine.base import _EngineTypingCommon
 from ...future import Connection
 from ...future import Engine
 from ...sql import Executable
 
 _TAsyncConnection = TypeVar("_TAsyncConnection", bound=AsyncConnection)
 _TAsyncTransaction = TypeVar("_TAsyncTransaction", bound=AsyncTransaction)
+_TEngine = TypeVar("_TEngine", bound=AsyncEngine)
 
 def create_async_engine(*arg: Any, **kw: Any) -> AsyncEngine: ...
 
 class AsyncConnectable: ...
 
 class AsyncConnection(
-    ProxyComparable, StartableContext["AsyncConnection"], AsyncConnectable
+    _ConnectionTypingCommon,
+    ProxyComparable,
+    StartableContext["AsyncConnection"],
+    AsyncConnectable,
 ):
     # copied from future.Connection via create_proxy_methods
     @property
     def closed(self) -> bool: ...
     @property
     def invalidated(self) -> bool: ...
-    dialect: Any
+    dialect: Dialect
     @property
     def default_isolation_level(self) -> Any: ...
     # end copied
@@ -99,12 +106,9 @@ class AsyncConnection(
         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 AsyncEngine(_EngineTypingCommon, ProxyComparable, AsyncConnectable):
+    @property
+    def engine(self: _TEngine) -> _TEngine: ...
     class _trans_ctx(StartableContext[AsyncConnection]):
         conn: AsyncConnection = ...
         def __init__(self, conn: AsyncConnection) -> None: ...

sqlalchemy-stubs/ext/asyncio/scoping.pyi

@@ -0,0 +1,19 @@
+from typing import Any
+from typing import Callable
+
+from .session import _AsyncSessionTypingCommon
+from .session import AsyncSession
+from ...orm.scoping import ScopedSessionMixin
+from ...util import ScopedRegistry
+
+class async_scoped_session(
+    _AsyncSessionTypingCommon, ScopedSessionMixin[AsyncSession]
+):
+    session_factory: Callable[..., AsyncSession] = ...
+    registry: ScopedRegistry = ...
+    def __init__(
+        self,
+        session_factory: Callable[..., AsyncSession],
+        scopefunc: Callable[..., Any] = ...,
+    ) -> None: ...
+    async def remove(self) -> None: ...