Improve future.* (#103)

Author: bryanforbes

sqlalchemy-stubs/future/__init__.pyi

@@ -1,7 +1,6 @@
-from typing import Any
-
 from .engine import Connection as Connection
 from .engine import create_engine as create_engine
 from .engine import Engine as Engine
+from ..sql.selectable import Select
 
-select: Any
+select = Select._create_future_select

sqlalchemy-stubs/future/engine.pyi

@@ -1,48 +1,98 @@
 from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import List
+from typing import NoReturn
 from typing import Optional
+from typing import Type
+from typing import Union
 
-from .. import util as util
-from ..engine import Connection as _LegacyConnection
-from ..engine import Engine as _LegacyEngine
-from ..engine.base import OptionEngineMixin as OptionEngineMixin
+from .. import util
+from ..engine.base import Connection as _LegacyConnection
+from ..engine.base import Engine as _LegacyEngine
+from ..engine.base import NestedTransaction
+from ..engine.base import OptionEngineMixin
+from ..engine.base import Transaction
+from ..engine.create import _Debug
+from ..engine.create import _IsolationLevel
+from ..engine.create import _ParamStyle
+from ..engine.create import _ResetOnReturn
+from ..engine.url import URL
+from ..pool import Pool
 
-NO_OPTIONS: Any
+NO_OPTIONS: util.immutabledict[Any, Any] = ...
 
-def create_engine(*arg: Any, **kw: Any): ...
+def create_engine(
+    url: Union[str, URL],
+    *,
+    case_sensitive: bool = ...,
+    connect_args: Dict[Any, Any] = ...,
+    convert_unicode: bool = ...,
+    creator: Callable[..., Any] = ...,
+    echo: Union[bool, _Debug] = ...,
+    echo_pool: Union[bool, _Debug] = ...,
+    enable_from_linting: bool = ...,
+    encoding: str = ...,
+    execution_options: Dict[Any, Any] = ...,
+    hide_parameters: bool = ...,
+    implicit_returning: bool = ...,
+    isolation_level: _IsolationLevel = ...,
+    json_deserializer: Callable[..., Any] = ...,
+    json_serializer: Callable[..., Any] = ...,
+    label_length: Optional[int] = ...,
+    listeners: Any = ...,
+    logging_name: str = ...,
+    max_identifier_length: Optional[int] = ...,
+    max_overflow: int = ...,
+    module: Optional[Any] = ...,
+    paramstyle: Optional[_ParamStyle] = ...,
+    pool: Optional[Pool] = ...,
+    poolclass: Optional[Type[Pool]] = ...,
+    pool_logging_name: str = ...,
+    pool_pre_ping: bool = ...,
+    pool_size: int = ...,
+    pool_recycle: int = ...,
+    pool_reset_on_return: Optional[_ResetOnReturn] = ...,
+    pool_timeout: float = ...,
+    pool_use_lifo: bool = ...,
+    plugins: List[str] = ...,
+    query_cache_size: int = ...,
+    **kw: Any,
+) -> Engine: ...
 
 class Connection(_LegacyConnection):
-    def begin(self): ...
-    def begin_nested(self): ...
+    def begin(self) -> Transaction: ...
+    def begin_nested(self) -> NestedTransaction: ...
     def commit(self) -> None: ...
     def rollback(self) -> None: ...
     def close(self) -> None: ...
-    def execute(
+    def execute(  # type: ignore[override]
         self,
         statement: Any,
         parameters: Optional[Any] = ...,
         execution_options: Optional[Any] = ...,
-    ): ...
-    def scalar(
+    ) -> Any: ...
+    def scalar(  # type: ignore[override]
         self,
         statement: Any,
         parameters: Optional[Any] = ...,
         execution_options: Optional[Any] = ...,
-    ): ...
+    ) -> Any: ...
 
 class Engine(_LegacyEngine):
-    transaction: Any = ...
-    run_callable: Any = ...
-    execute: Any = ...
-    scalar: Any = ...
-    table_names: Any = ...
-    has_table: Any = ...
+    def transaction(self, *arg: Any, **kw: Any) -> NoReturn: ...
+    def run_callable(self, *arg: Any, **kw: Any) -> NoReturn: ...
+    def execute(self, *arg: Any, **kw: Any) -> NoReturn: ...
+    def scalar(self, *arg: Any, **kw: Any) -> NoReturn: ...
+    def table_names(self, *arg: Any, **kw: Any) -> NoReturn: ...
+    def has_table(self, *arg: Any, **kw: Any) -> NoReturn: ...
     class _trans_ctx:
-        conn: Any = ...
-        def __init__(self, conn: Any) -> None: ...
-        transaction: Any = ...
-        def __enter__(self): ...
+        conn: Connection = ...
+        def __init__(self, conn: Connection) -> None: ...
+        transaction: Transaction = ...
+        def __enter__(self) -> Connection: ...
         def __exit__(self, type_: Any, value: Any, traceback: Any) -> None: ...
-    def begin(self): ...
-    def connect(self): ...
+    def begin(self) -> _trans_ctx: ...  # type: ignore[override]
+    def connect(self) -> Connection: ...  # type: ignore[override]
 
 class OptionEngine(OptionEngineMixin, Engine): ...

sqlalchemy-stubs/sql/selectable.pyi

@@ -532,6 +532,8 @@ class Select(
     ) -> _SEL: ...
     @classmethod
     def _create(cls, *args: Any, **kw: Any) -> Select: ...
+    @classmethod
+    def _create_future_select(cls: Type[_SEL], *entities: Any) -> _SEL: ...
     def __init__(self) -> None: ...
     def filter(self: _SEL, *criteria: Any) -> _SEL: ...
     def filter_by(self: _SEL, **kwargs: Any) -> _SEL: ...