Improve engine.base (#90)

* Improve `engine.base`

* Updates based on PR feedback

Author: bryanforbes

sqlalchemy-stubs/engine/base.pyi

@@ -1,185 +1,248 @@
 from typing import Any
+from typing import Dict
+from typing import List
 from typing import Optional
+from typing import TypeVar
+from typing import Union
 
+from typing_extensions import Literal
+from typing_extensions import Protocol
+
+from .cursor import CursorResult
+from .interfaces import _DBAPIConnection
+from .interfaces import _DBAPICursor
 from .interfaces import Connectable as Connectable
-from .interfaces import ExceptionContext as ExceptionContext
-from .. import exc as exc
-from .. import inspection as inspection
-from .. import log as log
-from .. import util as util
-from ..sql import compiler as compiler
+from .interfaces import Dialect
+from .interfaces import ExceptionContext
+from .interfaces import ExecutionContext
+from .url import URL
+from .. import log
+from .. import util
+from ..exc import StatementError
+from ..pool import Pool
+
+_T = TypeVar("_T")
+_T_co = TypeVar("_T_co", covariant=True)
+_T_contra = TypeVar("_T_contra", contravariant=True)
+_TConnection = TypeVar("_TConnection", bound=Connection)
+_TTransaction = TypeVar("_TTransaction", bound=Transaction)
+_TEngine = TypeVar("_TEngine", bound=Engine)
+
+_ExecutionOptions: util.immutabledict[Any, Any]
 
-_ExecutionOptions = Any
+class _ConnectionCallable(Protocol[_T_contra, _T_co]):
+    def __call__(
+        self, __connection: _T_contra, *args: Any, **kwargs: Any
+    ) -> _T_co: ...
 
 class Connection(Connectable):
-    engine: Any = ...
-    dialect: Any = ...
+    engine: Engine = ...
+    dialect: Dialect = ...
     should_close_with_result: bool = ...
     dispatch: Any = ...
     def __init__(
         self,
-        engine: Any,
-        connection: Optional[Any] = ...,
+        engine: Engine,
+        connection: Optional[_DBAPIConnection] = ...,
         close_with_result: bool = ...,
         _branch_from: Optional[Any] = ...,
-        _execution_options: Optional[Any] = ...,
+        _execution_options: Optional[Dict[str, Any]] = ...,
         _dispatch: Optional[Any] = ...,
         _has_events: Optional[Any] = ...,
     ) -> None: ...
-    def schema_for_object(self, obj: Any): ...
-    def __enter__(self): ...
+    def schema_for_object(self, obj: Any) -> str: ...
+    def __enter__(self: _TConnection) -> _TConnection: ...
     def __exit__(self, type_: Any, value: Any, traceback: Any) -> None: ...
-    def execution_options(self, **opt: Any): ...
-    def get_execution_options(self): ...
+    def execution_options(self: _TConnection, **opt: Any) -> _TConnection: ...
+    def get_execution_options(self) -> Dict[str, Any]: ...
     @property
-    def closed(self): ...
+    def closed(self) -> bool: ...
     @property
-    def invalidated(self): ...
+    def invalidated(self) -> bool: ...
     @property
-    def connection(self): ...
-    def get_isolation_level(self): ...
+    def connection(self) -> _DBAPIConnection: ...
+    def get_isolation_level(self) -> Any: ...
     @property
-    def default_isolation_level(self): ...
+    def default_isolation_level(self) -> Any: ...
     @property
-    def info(self): ...
-    def connect(self, close_with_result: bool = ...): ...
-    def invalidate(self, exception: Optional[Any] = ...): ...
+    def info(self) -> Dict[Any, Any]: ...
+    def connect(self: _TConnection, close_with_result: bool = ...) -> _TConnection: ...  # type: ignore[override]
+    def invalidate(self, exception: Optional[Any] = ...) -> None: ...
     def detach(self) -> None: ...
-    def begin(self): ...
-    def begin_nested(self): ...
-    def begin_twophase(self, xid: Optional[Any] = ...): ...
-    def recover_twophase(self): ...
-    def rollback_prepared(self, xid: Any, recover: bool = ...) -> None: ...
-    def commit_prepared(self, xid: Any, recover: bool = ...) -> None: ...
-    def in_transaction(self): ...
-    def in_nested_transaction(self): ...
-    def get_transaction(self): ...
-    def get_nested_transaction(self): ...
+    def begin(self) -> Optional[Transaction]: ...
+    def begin_nested(self) -> NestedTransaction: ...
+    def begin_twophase(
+        self, xid: Optional[str] = ...
+    ) -> TwoPhaseTransaction: ...
+    def recover_twophase(self) -> None: ...
+    def rollback_prepared(self, xid: str, recover: bool = ...) -> None: ...
+    def commit_prepared(self, xid: str, recover: bool = ...) -> None: ...
+    def in_transaction(self) -> bool: ...
+    def in_nested_transaction(self) -> bool: ...
+    def get_transaction(self) -> Optional[Transaction]: ...
+    def get_nested_transaction(self) -> Optional[NestedTransaction]: ...
     def close(self) -> None: ...
-    def scalar(self, object_: Any, *multiparams: Any, **params: Any): ...
-    def execute(self, statement: Any, *multiparams: Any, **params: Any): ...
+    def scalar(
+        self, object_: Any, *multiparams: Any, **params: Any
+    ) -> Any: ...
+    def execute(self, statement: Any, *multiparams: Any, **params: Any) -> CursorResult: ...  # type: ignore[override]
     def exec_driver_sql(
         self,
-        statement: Any,
+        statement: str,
         parameters: Optional[Any] = ...,
         execution_options: Optional[Any] = ...,
-    ): ...
-    def transaction(self, callable_: Any, *args: Any, **kwargs: Any): ...
-    def run_callable(self, callable_: Any, *args: Any, **kwargs: Any): ...
+    ) -> CursorResult: ...
+    def transaction(
+        self: _TConnection,
+        callable_: _ConnectionCallable[_TConnection, _T],
+        *args: Any,
+        **kwargs: Any,
+    ) -> _T: ...
+    def run_callable(
+        self: _TConnection,
+        callable_: _ConnectionCallable[_TConnection, _T],
+        *args: Any,
+        **kwargs: Any,
+    ) -> _T: ...
 
 class ExceptionContextImpl(ExceptionContext):
-    engine: Any = ...
-    connection: Any = ...
-    sqlalchemy_exception: Any = ...
-    original_exception: Any = ...
-    execution_context: Any = ...
-    statement: Any = ...
-    parameters: Any = ...
-    is_disconnect: Any = ...
-    invalidate_pool_on_disconnect: Any = ...
+    engine: Engine = ...
+    connection: Connection = ...
+    sqlalchemy_exception: Optional[StatementError] = ...
+    original_exception: BaseException = ...
+    execution_context: Optional[ExecutionContext] = ...
+    statement: Optional[str] = ...
+    parameters: Optional[Any] = ...
+    is_disconnect: bool = ...
+    invalidate_pool_on_disconnect: bool = ...
     def __init__(
         self,
-        exception: Any,
-        sqlalchemy_exception: Any,
-        engine: Any,
-        connection: Any,
-        cursor: Any,
-        statement: Any,
-        parameters: Any,
-        context: Any,
-        is_disconnect: Any,
-        invalidate_pool_on_disconnect: Any,
+        exception: BaseException,
+        sqlalchemy_exception: Optional[StatementError],
+        engine: Engine,
+        connection: Optional[Connection],
+        cursor: Optional[_DBAPICursor],
+        statement: Optional[str],
+        parameters: Optional[Any],
+        context: Optional[ExecutionContext],
+        is_disconnect: bool,
+        invalidate_pool_on_disconnect: bool,
     ) -> None: ...
 
 class Transaction:
-    def __init__(self, connection: Any) -> None: ...
+    def __init__(self, connection: Connection) -> None: ...
     @property
-    def is_valid(self): ...
+    def is_valid(self) -> bool: ...
     def close(self) -> None: ...
     def rollback(self) -> None: ...
     def commit(self) -> None: ...
-    def __enter__(self): ...
+    def __enter__(self: _TTransaction) -> _TTransaction: ...
     def __exit__(self, type_: Any, value: Any, traceback: Any) -> None: ...
 
 class MarkerTransaction(Transaction):
-    connection: Any = ...
-    def __init__(self, connection: Any) -> None: ...
+    connection: Connection = ...
+    def __init__(self, connection: Connection) -> None: ...
     @property
-    def is_active(self): ...
+    def is_active(self) -> bool: ...
 
 class RootTransaction(Transaction):
-    connection: Any = ...
+    connection: Connection = ...
     is_active: bool = ...
-    def __init__(self, connection: Any) -> None: ...
+    def __init__(self, connection: Connection) -> None: ...
 
 class NestedTransaction(Transaction):
-    connection: Any = ...
+    connection: Connection = ...
     is_active: bool = ...
-    def __init__(self, connection: Any) -> None: ...
+    def __init__(self, connection: Connection) -> None: ...
 
 class TwoPhaseTransaction(RootTransaction):
-    xid: Any = ...
-    def __init__(self, connection: Any, xid: Any) -> None: ...
+    xid: str = ...
+    def __init__(self, connection: Connection, xid: str) -> None: ...
     def prepare(self) -> None: ...
 
 class Engine(Connectable, log.Identified):
-    pool: Any = ...
-    url: Any = ...
-    dialect: Any = ...
-    logging_name: Any = ...
-    echo: Any = ...
-    hide_parameters: Any = ...
+    pool: Pool = ...
+    url: URL = ...
+    dialect: Dialect = ...
+    logging_name: Optional[str] = ...
+    echo: Optional[Union[bool, Literal["debug"]]] = ...
+    hide_parameters: bool = ...
     def __init__(
         self,
-        pool: Any,
-        dialect: Any,
-        url: Any,
-        logging_name: Optional[Any] = ...,
-        echo: Optional[Any] = ...,
+        pool: Pool,
+        dialect: Dialect,
+        url: URL,
+        logging_name: Optional[str] = ...,
+        echo: Optional[Union[bool, Literal["debug"]]] = ...,
         query_cache_size: int = ...,
-        execution_options: Optional[Any] = ...,
+        execution_options: Optional[Dict[str, Any]] = ...,
         hide_parameters: bool = ...,
     ) -> None: ...
     @property
-    def engine(self): ...
+    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): ...
-    def get_execution_options(self): ...
+    def execution_options(self, **opt: Any) -> OptionEngine: ...
+    def get_execution_options(self) -> Dict[str, Any]: ...
     @property
-    def name(self): ...
+    def name(self) -> str: ...
     @property
-    def driver(self): ...
+    def driver(self) -> str: ...
     def dispose(self) -> None: ...
     class _trans_ctx:
-        conn: Any = ...
-        transaction: Any = ...
-        close_with_result: Any = ...
+        conn: Connection = ...
+        transaction: Transaction = ...
+        close_with_result: bool = ...
         def __init__(
-            self, conn: Any, transaction: Any, close_with_result: Any
+            self,
+            conn: Connection,
+            transaction: Transaction,
+            close_with_result: bool,
         ) -> None: ...
-        def __enter__(self): ...
+        def __enter__(self) -> Connection: ...
         def __exit__(self, type_: Any, value: Any, traceback: Any) -> None: ...
-    def begin(self, close_with_result: bool = ...): ...
-    def transaction(self, callable_: Any, *args: Any, **kwargs: Any): ...
-    def run_callable(self, callable_: Any, *args: Any, **kwargs: Any): ...
-    def execute(self, statement: Any, *multiparams: Any, **params: Any): ...
-    def scalar(self, statement: Any, *multiparams: Any, **params: Any): ...
-    def connect(self, close_with_result: bool = ...): ...
+    def begin(self, close_with_result: bool = ...) -> _trans_ctx: ...
+    def transaction(
+        self,
+        callable_: _ConnectionCallable[_TConnection, _T],
+        *args: Any,
+        **kwargs: Any,
+    ) -> _T: ...
+    def run_callable(
+        self,
+        callable_: _ConnectionCallable[_TConnection, _T],
+        *args: Any,
+        **kwargs: Any,
+    ) -> _T: ...
+    def execute(  # type: ignore[override]
+        self, statement: Any, *multiparams: Any, **params: Any
+    ) -> CursorResult: ...
+    def scalar(  # type: ignore[override]
+        self, statement: Any, *multiparams: Any, **params: Any
+    ) -> Any: ...
+    def connect(self, close_with_result: bool = ...) -> Connection: ...  # type: ignore[override]
     def table_names(
-        self, schema: Optional[Any] = ..., connection: Optional[Any] = ...
-    ): ...
-    def has_table(self, table_name: Any, schema: Optional[Any] = ...): ...
-    def raw_connection(self, _connection: Optional[Any] = ...): ...
+        self,
+        schema: Optional[str] = ...,
+        connection: Optional[Connection] = ...,
+    ) -> List[str]: ...
+    def has_table(
+        self, table_name: str, schema: Optional[str] = ...
+    ) -> bool: ...
+    def raw_connection(
+        self, _connection: Optional[Connection] = ...
+    ) -> _DBAPIConnection: ...
 
 class OptionEngineMixin:
-    url: Any = ...
-    dialect: Any = ...
-    logging_name: Any = ...
-    echo: Any = ...
-    hide_parameters: Any = ...
+    url: URL = ...
+    dialect: Dialect = ...
+    logging_name: Optional[str] = ...
+    echo: Optional[Union[bool, Literal["debug"]]] = ...
+    hide_parameters: bool = ...
     dispatch: Any = ...
-    def __init__(self, proxied: Any, execution_options: Any) -> None: ...
-    pool: Any = ...
+    pool: Pool = ...
+    def __init__(
+        self, proxied: Engine, execution_options: Dict[str, Any]
+    ) -> None: ...
 
 class OptionEngine(OptionEngineMixin, Engine): ...