Merge pull request #25 from bryanforbes/improve-sql-base

Improve `sql.base` typings

Author: zzzeek

sqlalchemy-stubs/sql/base.pyi

@@ -1,97 +1,120 @@
 from types import ModuleType
 from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import Generic
+from typing import Iterable
+from typing import Iterator
+from typing import List
+from typing import Mapping
+from typing import NoReturn
 from typing import Optional
-
-from . import roles as roles
-from .traversals import HasCacheKey as HasCacheKey
-from .traversals import HasCopyInternals as HasCopyInternals
-from .traversals import MemoizedHasCacheKey as MemoizedHasCacheKey
-from .visitors import ClauseVisitor as ClauseVisitor
-from .visitors import ExtendedInternalTraversal as ExtendedInternalTraversal
-from .visitors import InternalTraversal as InternalTraversal
-from .. import exc as exc
-from .. import util as util
+from typing import overload
+from typing import Tuple
+from typing import Type
+from typing import TypeVar
+from typing import Union
+
+from . import roles
+from .schema import Column
+from .traversals import HasCacheKey
+from .traversals import HasCopyInternals
+from .visitors import ClauseVisitor
+from .. import util
+from ..engine import Connection
+from ..engine import Engine
 from ..util import HasMemoized as HasMemoized
-from ..util import hybridmethod as hybridmethod
+from ..util import langhelpers
+
+_T = TypeVar("_T")
+_SC = TypeVar("_SC", bound=SingletonConstant)
+_O = TypeVar("_O", bound=Options)
+_E = TypeVar("_E", bound=Executable)
+_C = TypeVar("_C", bound=Column[Any])
+_OC = TypeVar("_OC", bound=Column[Any])
 
 coercions: ModuleType
 elements: ModuleType
 type_api: ModuleType
-PARSE_AUTOCOMMIT: Any
-NO_ARG: Any
+PARSE_AUTOCOMMIT: langhelpers._symbol
+NO_ARG: langhelpers._symbol
 
 class Immutable:
-    def unique_params(self, *optionaldict: Any, **kwargs: Any) -> None: ...
-    def params(self, *optionaldict: Any, **kwargs: Any) -> None: ...
+    def unique_params(self, *optionaldict: Any, **kwargs: Any) -> NoReturn: ...
+    def params(self, *optionaldict: Any, **kwargs: Any) -> NoReturn: ...
 
 class SingletonConstant(Immutable):
-    def __new__(cls, *arg: Any, **kw: Any): ...
+    def __new__(cls: Type[_SC], *arg: Any, **kw: Any) -> _SC: ...
 
-class _DialectArgView(util.collections_abc.MutableMapping):
+class _DialectArgView(util.collections_abc.MutableMapping[str, Any]):
     obj: Any = ...
     def __init__(self, obj: Any) -> None: ...
-    def __getitem__(self, key: Any): ...
-    def __setitem__(self, key: Any, value: Any) -> None: ...
-    def __delitem__(self, key: Any) -> None: ...
-    def __len__(self): ...
-    def __iter__(self) -> Any: ...
+    def __getitem__(self, key: str) -> Any: ...
+    def __setitem__(self, key: str, value: Any) -> None: ...
+    def __delitem__(self, key: str) -> None: ...
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[str]: ...
 
-class _DialectArgDict(util.collections_abc.MutableMapping):
+class _DialectArgDict(util.collections_abc.MutableMapping[str, Any]):
     def __init__(self) -> None: ...
-    def __len__(self): ...
-    def __iter__(self) -> Any: ...
-    def __getitem__(self, key: Any): ...
-    def __setitem__(self, key: Any, value: Any) -> None: ...
-    def __delitem__(self, key: Any) -> None: ...
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[str]: ...
+    def __getitem__(self, key: str) -> Any: ...
+    def __setitem__(self, key: str, value: Any) -> None: ...
+    def __delitem__(self, key: str) -> None: ...
 
 class DialectKWArgs:
     @classmethod
     def argument_for(
         cls, dialect_name: Any, argument_name: Any, default: Any
     ) -> None: ...
-    def dialect_kwargs(self): ...
+    @util.memoized_property
+    def dialect_kwargs(self) -> _DialectArgView: ...
     @property
-    def kwargs(self): ...
-    def dialect_options(self): ...
+    def kwargs(self) -> _DialectArgView: ...
+    @util.memoized_property
+    def dialect_options(self) -> util.PopulateDict[str, _DialectArgDict]: ...
 
 class CompileState:
     plugins: Any = ...
     @classmethod
     def create_for_statement(
         cls, statement: Any, compiler: Any, **kw: Any
-    ): ...
+    ) -> Any: ...
     statement: Any = ...
     def __init__(self, statement: Any, compiler: Any, **kw: Any) -> None: ...
     @classmethod
-    def get_plugin_class(cls, statement: Any): ...
+    def get_plugin_class(cls, statement: Any) -> Optional[Any]: ...
     @classmethod
-    def plugin_for(cls, plugin_name: Any, visit_name: Any): ...
+    def plugin_for(
+        cls, plugin_name: Any, visit_name: Any
+    ) -> Callable[[_T], _T]: ...
 
 class Generative(HasMemoized): ...
 class InPlaceGenerative(HasMemoized): ...
 class HasCompileState(Generative): ...
 
 class _MetaOptions(type):
     def __init__(cls, classname: Any, bases: Any, dict_: Any) -> None: ...
-    def __add__(self, other: Any): ...
+    def __add__(self: Type[_T], other: Any) -> _T: ...
 
-class Options:
+class Options(metaclass=_MetaOptions):
     def __init__(self, **kw: Any) -> None: ...
-    def __add__(self, other: Any): ...
-    def __eq__(self, other: Any) -> Any: ...
+    def __add__(self: _O, other: Any) -> _O: ...
+    def __eq__(self, other: Any) -> bool: ...
     @classmethod
-    def isinstance(cls, klass: Any): ...
-    def add_to_element(self, name: Any, value: Any): ...
+    def isinstance(cls, klass: Any) -> bool: ...
+    def add_to_element(self: _O, name: Any, value: Any) -> _O: ...
     @classmethod
-    def safe_merge(cls, other: Any): ...
+    def safe_merge(cls: Type[_O], other: Any) -> _O: ...
     @classmethod
     def from_execution_options(
         cls,
         key: Any,
         attrs: Any,
         exec_options: Any,
         statement_exec_options: Any,
-    ): ...
+    ) -> Tuple[Any, Any]: ...
 
 class CacheableOptions(Options, HasCacheKey): ...
 
@@ -106,63 +129,72 @@ class Executable(roles.CoerceTextStatementRole, Generative):
     is_text: bool = ...
     is_delete: bool = ...
     is_dml: bool = ...
-    def options(self, *options: Any) -> None: ...
-    def execution_options(self, **kw: Any) -> None: ...
-    def get_execution_options(self): ...
-    def execute(self, *multiparams: Any, **params: Any): ...
-    def scalar(self, *multiparams: Any, **params: Any): ...
+    def options(self: _E, *options: Any) -> _E: ...
+    def execution_options(self: _E, **kw: Any) -> _E: ...
+    def get_execution_options(self) -> Any: ...
+    def execute(self, *multiparams: Any, **params: Any) -> Any: ...
+    def scalar(self, *multiparams: Any, **params: Any) -> Any: ...
     @property
-    def bind(self): ...
+    def bind(self) -> Optional[Union[Engine, Connection]]: ...
 
-class prefix_anon_map(dict):
-    def __missing__(self, key: Any): ...
+class prefix_anon_map(Dict[str, str]):
+    def __missing__(self, key: str) -> str: ...
 
 class SchemaEventTarget: ...
 
 class SchemaVisitor(ClauseVisitor):
     __traverse_options__: Any = ...
 
-class ColumnCollection:
-    def __init__(self, columns: Optional[Any] = ...) -> None: ...
-    def keys(self): ...
-    def __bool__(self): ...
-    def __len__(self): ...
-    def __iter__(self) -> Any: ...
-    def __getitem__(self, key: Any): ...
-    def __getattr__(self, key: Any): ...
-    def __contains__(self, key: Any): ...
-    def compare(self, other: Any): ...
-    def __eq__(self, other: Any) -> Any: ...
-    def get(self, key: Any, default: Optional[Any] = ...): ...
-    def __setitem__(self, key: Any, value: Any) -> None: ...
-    def __delitem__(self, key: Any) -> None: ...
-    def __setattr__(self, key: Any, obj: Any) -> None: ...
+class ColumnCollection(Generic[_C]):
+    def __init__(
+        self, columns: Optional[Iterable[Tuple[str, _C]]] = ...
+    ) -> None: ...
+    def keys(self) -> List[_C]: ...
+    def values(self) -> List[str]: ...
+    def items(self) -> List[Tuple[str, _C]]: ...
+    def __bool__(self) -> bool: ...
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[_C]: ...
+    def __getitem__(self, key: str) -> _C: ...
+    def __getattr__(self, key: str) -> _C: ...
+    def __contains__(self, key: str) -> bool: ...
+    def compare(self, other: Any) -> bool: ...
+    def __eq__(self, other: Any) -> bool: ...
+    @overload
+    def get(self, key: str) -> Optional[_C]: ...
+    @overload
+    def get(self, key: str, default: _T) -> Union[_C, _T]: ...
+    def __setitem__(self, key: str, value: _C) -> None: ...
+    def __delitem__(self, key: str) -> None: ...
+    def __setattr__(self, key: str, obj: _C) -> None: ...
     def clear(self) -> None: ...
-    def remove(self, column: Any) -> None: ...
-    def update(self, iter_: Any) -> None: ...
+    def remove(self, column: _C) -> None: ...
+    def update(
+        self, iter_: Union[Mapping[str, _C], Iterable[Tuple[str, _C]]]
+    ) -> None: ...
     __hash__: Any = ...
-    def add(self, column: Any, key: Optional[Any] = ...) -> None: ...
-    def contains_column(self, col: Any): ...
-    def as_immutable(self): ...
+    def add(self, column: _C, key: Optional[str] = ...) -> None: ...
+    def contains_column(self, col: Column[Any]) -> bool: ...
+    def as_immutable(self) -> ImmutableColumnCollection[_C]: ...
     def corresponding_column(
-        self, column: Any, require_embedded: bool = ...
-    ): ...
-
-class DedupeColumnCollection(ColumnCollection):
-    def add(self, column: Any, key: Optional[Any] = ...) -> None: ...
-    def extend(self, iter_: Any) -> None: ...
-    def remove(self, column: Any) -> None: ...
-    def replace(self, column: Any) -> None: ...
-
-class ImmutableColumnCollection(util.ImmutableContainer, ColumnCollection):
-    def __init__(self, collection: Any) -> None: ...
-    add: Any = ...
-    extend: Any = ...
-    remove: Any = ...
-
-class ColumnSet(util.ordered_column_set):
-    def contains_column(self, col: Any): ...
-    def extend(self, cols: Any) -> None: ...
-    def __add__(self, other: Any): ...
+        self, column: Column[Any], require_embedded: bool = ...
+    ) -> Optional[_C]: ...
+
+class DedupeColumnCollection(ColumnCollection[_C]):
+    def add(self, column: _C, key: Optional[str] = ...) -> None: ...
+    def extend(self, iter_: Iterable[_C]) -> None: ...
+    def remove(self, column: _C) -> None: ...
+    def replace(self, column: _C) -> None: ...
+
+class ImmutableColumnCollection(util.ImmutableContainer, ColumnCollection[_C]):
+    def __init__(self, collection: ColumnCollection[_C]) -> None: ...
+    def add(self, column: _C, key: Optional[str] = ...) -> NoReturn: ...
+    def extend(self, iter_: Iterable[_C]) -> NoReturn: ...
+    def remove(self, column: _C) -> NoReturn: ...
+
+class ColumnSet(util.ordered_column_set[_C]):
+    def contains_column(self, col: Column[Any]) -> bool: ...
+    def extend(self, cols: Iterable[_C]) -> None: ...
+    def __add__(self, other: Iterable[_OC]) -> List[Union[_C, _OC]]: ...  # type: ignore[override]
     def __eq__(self, other: Any) -> Any: ...
-    def __hash__(self) -> Any: ...
+    def __hash__(self) -> int: ...  # type: ignore[override]