Various updates to sqlalchemy.sql typings (#46)

Author: bryanforbes

sqlalchemy-stubs/sql/annotation.pyi

@@ -1,8 +1,11 @@
-from typing import Any, Dict
+from typing import Any, Dict, Type, TypeVar, Generic
 
-class Annotated(object):
+_T = TypeVar('_T')
+
+class Annotated(Generic[_T]):
+    __element: _T
     def __new__(cls, *args): ...
-    def __init__(self, element, values) -> None: ...
+    def __init__(self, element: _T, values: Any) -> None: ...
     def __hash__(self): ...
     def __eq__(self, other): ...
     def _annotate(self, values): ...
@@ -12,4 +15,4 @@ class Annotated(object):
     def _constructor(self): ...
     def _clone(self): ...
 
-annotated_classes: Dict[Any, Any] = ...
+annotated_classes: Dict[Type[Any], Type[Annotated[Any]]] = ...

sqlalchemy-stubs/sql/base.pyi

@@ -14,8 +14,8 @@ NO_ARG: util.symbol = ...
 _I = TypeVar('_I', bound=Immutable)
 
 class Immutable(object):
-    def unique_params(self, *optionaldict: Any, **kwargs: Any): ...
-    def params(self, *optionaldict: Any, **kwargs: Any): ...
+    def unique_params(self, *optionaldict: Any, **kwargs: Any) -> Any: ...
+    def params(self, *optionaldict: Any, **kwargs: Any) -> Any: ...
     def _clone(self: _I) -> _I: ...
 
 class _DialectArgView(MutableMapping[str, Any]):

sqlalchemy-stubs/sql/compiler.pyi

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Set, Dict, List, Set, Pattern, TypeVar, Type
+from typing import Any, Optional, Union, Set, Dict, List, Set, Pattern, TypeVar, Type, Mapping, Tuple, Callable
 from .. import util, engine
 from .elements import ClauseElement, BindParameter
 from .type_api import TypeEngine
@@ -17,21 +17,22 @@ COMPOUND_KEYWORDS: Dict[Any, str] = ...
 class Compiled(object):
     execution_options: Any = ...
     dialect: engine.Dialect = ...
-    bind: Any = ...
-    preparer: Any = ...
+    bind: Optional[Union[engine.Engine, engine.Connection]] = ...
+    preparer: Optional[IdentifierPreparer] = ...
     statement: ClauseElement = ...
     can_execute: Optional[bool] = ...
     string: Optional[str] = ...
-    def __init__(self, dialect: engine.Dialect, statement: ClauseElement, bind: Optional[Any] = ...,
-                 schema_translate_map: Optional[Any] = ..., compile_kwargs: Any = ...) -> None: ...
+    def __init__(self, dialect: engine.Dialect, statement: ClauseElement,
+                 bind: Optional[Union[engine.Engine, engine.Connection]] = ...,
+                 schema_translate_map: Optional[Any] = ..., compile_kwargs: Mapping[Any, Any] = ...) -> None: ...
     def compile(self) -> str: ...
     @property
-    def sql_compiler(self) -> Any: ...
+    def sql_compiler(self) -> SQLCompiler: ...
     def process(self, obj: Any, **kwargs: Any) -> str: ...
     def __str__(self) -> str: ...
-    def construct_params(self, params: Optional[Dict[str, Any]] = ...) -> Any: ...
+    def construct_params(self, params: Optional[Dict[str, Any]] = ...) -> Dict[str, Any]: ...
     @property
-    def params(self) -> Optional[Dict[str, Any]]: ...
+    def params(self) -> Dict[str, Any]: ...
     def execute(self, *multiparams: Any, **params: Any) -> engine.ResultProxy: ...
     def scalar(self, *multiparams: Any, **params: Any) -> Any: ...
 
@@ -50,35 +51,35 @@ class SQLCompiler(Compiled):
     isinsert: bool = ...
     isupdate: bool = ...
     isplaintext: bool = ...
-    returning: Any = ...
+    returning: bool = ...
     returning_precedes_values: bool = ...
     render_table_with_column_in_update_from: bool = ...
     ansi_bind_rules: bool = ...
-    insert_prefetch: Any = ...
-    update_prefetch: Any = ...
+    insert_prefetch: Tuple[Any, ...] = ...
+    update_prefetch: Tuple[Any, ...] = ...
     column_keys: Optional[List[str]] = ...
     inline: bool = ...
     binds: Dict[str, BindParameter[Any]] = ...
-    bind_names: Any = ...
-    stack: Any = ...
-    positional: Any = ...
-    positiontup: Any = ...
-    bindtemplate: Any = ...
-    ctes: Any = ...
-    label_length: Any = ...
-    anon_map: Any = ...
-    truncated_names: Any = ...
+    bind_names: Dict[BindParameter[Any], str] = ...
+    stack: List[Dict[str, Any]] = ...
+    positional: bool = ...
+    positiontup: List[Any] = ...
+    bindtemplate: str = ...
+    ctes: Optional[util.OrderedDict[str, Any]] = ...
+    label_length: int = ...
+    anon_map: util.PopulateDict = ...
+    truncated_names: Dict[str, Any] = ...
     def __init__(self, dialect: engine.Dialect, statement: ClauseElement, column_keys: Optional[List[str]] = ...,
                  inline: bool = ..., **kwargs: Any) -> None: ...
     @property
     def prefetch(self) -> List[Any]: ...
     def is_subquery(self) -> bool: ...
     @property
     def sql_compiler(self: _SQLC) -> _SQLC: ...
-    def construct_params(self, params: Optional[Any] = ..., _group_number: Optional[Any] = ...,
-                         _check: bool = ...) -> Optional[Dict[str, Any]]: ...
+    def construct_params(self, params: Optional[Dict[str, Any]] = ..., _group_number: Optional[Any] = ...,
+                         _check: bool = ...) -> Dict[str, Any]: ...
     @property
-    def params(self) -> Optional[Dict[str, Any]]: ...
+    def params(self) -> Dict[str, Any]: ...
     def default_from(self) -> str: ...
     def visit_grouping(self, grouping: Any, asfrom: bool = ..., **kwargs: Any) -> str: ...
     def visit_label_reference(self, element, within_columns_clause: bool = ..., **kwargs: Any) -> str: ...
@@ -178,10 +179,12 @@ class SQLCompiler(Compiled):
     def visit_release_savepoint(self, savepoint_stmt) -> str: ...
 
 class StrSQLCompiler(SQLCompiler):
-    def visit_getitem_binary(self, binary, operator, **kw): ...
-    def visit_json_getitem_op_binary(self, binary, operator, **kw): ...
-    def visit_json_path_getitem_op_binary(self, binary, operator, **kw): ...
-    def returning_clause(self, stmt, returning_cols): ...
+    def visit_getitem_binary(self, binary, operator, **kw: Any) -> str: ...
+    def visit_json_getitem_op_binary(self, binary, operator, **kw: Any) -> str: ...
+    def visit_json_path_getitem_op_binary(self, binary, operator, **kw: Any) -> str: ...
+    def returning_clause(self, stmt, returning_cols) -> str: ...
+    def update_from_clause(self, update_stmt, from_table, extra_froms, from_hints, **kw: Any) -> str: ...
+    def delete_extra_from_clause(self, update_stmt, from_table, extra_froms, from_hints, **kw: Any) -> str: ...
 
 class DDLCompiler(Compiled):
     @property
@@ -271,7 +274,7 @@ class IdentifierPreparer(object):
     reserved_words: Set[str] = ...
     legal_characters: Pattern[str] = ...
     illegal_initial_characters: Pattern[str] = ...
-    schema_for_object: Any = ...
+    schema_for_object: Callable[[Any], str] = ...
     dialect: engine.Dialect = ...
     initial_quote: str = ...
     final_quote: str = ...

sqlalchemy-stubs/sql/crud.pyi

@@ -1,6 +1,6 @@
-from typing import Any
+from .. import util
 
-REQUIRED: Any = ...
-ISINSERT: Any = ...
-ISUPDATE: Any = ...
-ISDELETE: Any = ...
+REQUIRED: util.symbol = ...
+ISINSERT: util.symbol = ...
+ISUPDATE: util.symbol = ...
+ISDELETE: util.symbol = ...

sqlalchemy-stubs/sql/ddl.pyi

@@ -1,55 +1,55 @@
-from typing import Any, Optional, Union, List, Dict, TypeVar, Generic, Sequence, Callable, Tuple
+from typing import Any, Optional, Union, List, Dict, TypeVar, Generic, Sequence, Callable, Tuple, Set
 from typing_extensions import Protocol
 from .elements import ClauseElement as ClauseElement, ColumnElement
 from .base import Executable as Executable, SchemaVisitor as SchemaVisitor
 from .schema import Table, MetaData, SchemaItem, ForeignKeyConstraint, ForeignKey
 from .compiler import IdentifierPreparer
 from .type_api import TypeEngine
-from .. import engine
+from ..engine import Dialect, Engine, Connection, ResultProxy, Connectable
 
 _T = TypeVar('_T')
 
 class _DDLCompiles(ClauseElement): ...
 
 class _DDLCallable(Protocol):
-    def __call__(self, ddl: DDLElement, target: Optional[Union[Table, MetaData]], bind: engine.Connection,
+    def __call__(self, ddl: DDLElement, target: Optional[Union[Table, MetaData]], bind: Connection,
                  tables: Optional[List[Any]] = ..., state: Optional[Any] = ...,
                  checkfirst: bool = ...) -> bool: ...
 
 class _DDLOnCallback(Protocol):
     def __call__(self, ddl: DDLElement, event: Optional[str], target: Optional[Union[Table, MetaData]],
-                 connection: engine.Connection, tables: Optional[List[Any]] = ...) -> bool: ...
+                 connection: Connection, tables: Optional[List[Any]] = ...) -> bool: ...
 
 _DDLE = TypeVar('_DDLE', bound=DDLElement)
 
 class DDLElement(Executable, _DDLCompiles):
     target: Optional[SchemaItem] = ...
     on: Optional[Union[str, Tuple[str, ...], _DDLOnCallback]] = ...
-    dialect: Optional[engine.Dialect] = ...
+    dialect: Optional[Dialect] = ...
     callable_: Optional[_DDLCallable] = ...
-    bind: Optional[engine.Connection] = ...
-    def execute(self, bind: Optional[engine.Connectable] = ...,  # type: ignore  # incompatible with Executable.execute
-                target: Optional[SchemaItem] = ...) -> Optional[engine.ResultProxy]: ...
+    bind: Optional[Union[Engine, Connection]] = ...
+    def execute(self, bind: Optional[Union[Engine, Connection]] = ...,  # type: ignore  # incompatible with Executable.execute
+                target: Optional[SchemaItem] = ...) -> Optional[ResultProxy]: ...
     def execute_at(self, event_name: str, target: SchemaItem) -> None: ...
     def against(self: _DDLE, target: SchemaItem) -> _DDLE: ...
     state: Any = ...
-    def execute_if(self: _DDLE, dialect: Optional[engine.Dialect] = ..., callable_: Optional[_DDLCallable] = ...,
+    def execute_if(self: _DDLE, dialect: Optional[Dialect] = ..., callable_: Optional[_DDLCallable] = ...,
                    state: Optional[Any] = ...) -> _DDLE: ...
-    def __call__(self, target, bind, **kw): ...
+    def __call__(self, target: Optional[Union[Table, MetaData]], bind: Connection, **kw: Any) -> Optional[ResultProxy]: ...
 
 class DDL(DDLElement):
     __visit_name__: str = ...
     statement: str = ...
     context: Dict[Any, Any] = ...
     on: Optional[Union[str, Tuple[str, ...], _DDLOnCallback]] = ...
     def __init__(self, statement: str, on: Optional[Union[str, Tuple[str, ...], _DDLOnCallback]] = ...,
-                 context: Optional[Dict[Any, Any]] = ..., bind: Optional[engine.Connectable] = ...) -> None: ...
+                 context: Optional[Dict[Any, Any]] = ..., bind: Optional[Connectable] = ...) -> None: ...
 
 class _CreateDropBase(DDLElement, Generic[_T]):
     element: _T = ...
-    bind: Optional[engine.Connection] = ...
+    bind: Optional[Connection] = ...
     def __init__(self, element: _T, on: Optional[Union[str, Tuple[str, ...], _DDLOnCallback]] = ...,
-                 bind: Optional[engine.Connection] = ...) -> None: ...
+                 bind: Optional[Connection] = ...) -> None: ...
 
 class CreateSchema(_CreateDropBase[str]):
     __visit_name__: str = ...
@@ -67,7 +67,7 @@ class CreateTable(_CreateDropBase[Table]):
     columns: List[CreateColumn[Any]] = ...
     include_foreign_key_constraints: Optional[Sequence[ForeignKeyConstraint]] = ...
     def __init__(self, element: Table, on: Optional[Union[str, Tuple[str, ...], _DDLOnCallback]] = ...,
-                 bind: Optional[engine.Connectable] = ...,
+                 bind: Optional[Connectable] = ...,
                  include_foreign_key_constraints: Optional[Sequence[ForeignKeyConstraint]] = ...) -> None: ...
 
 class _DropView(_CreateDropBase[str]):
@@ -100,19 +100,19 @@ class AddConstraint(_CreateDropBase[str]):
 class DropConstraint(_CreateDropBase[str]):
     __visit_name__: str = ...
     cascade: bool = ...
-    def __init__(self, element: str, cascade: bool = ..., **kw) -> None: ...
+    def __init__(self, element: str, cascade: bool = ..., **kw: Any) -> None: ...
 
 class DDLBase(SchemaVisitor):
-    connection: engine.Connection = ...
-    def __init__(self, connection: engine.Connection) -> None: ...
+    connection: Connection = ...
+    def __init__(self, connection: Connection) -> None: ...
 
 class SchemaGenerator(DDLBase):
     checkfirst: bool = ...
     tables: Optional[List[Table]] = ...
     preparer: IdentifierPreparer = ...
-    dialect: engine.Dialect = ...
+    dialect: Dialect = ...
     memo: Dict[Any, Any] = ...
-    def __init__(self, dialect: engine.Dialect, connection: engine.Connection, checkfirst: bool = ...,
+    def __init__(self, dialect: Dialect, connection: Connection, checkfirst: bool = ...,
                  tables: Optional[List[Table]] = ..., **kwargs: Any) -> None: ...
     def visit_metadata(self, metadata: MetaData) -> None: ...
     def visit_table(self, table: Table, create_ok: bool = ...,
@@ -126,9 +126,9 @@ class SchemaDropper(DDLBase):
     checkfirst: bool = ...
     tables: Optional[List[Table]] = ...
     preparer: IdentifierPreparer = ...
-    dialect: engine.Dialect = ...
+    dialect: Dialect = ...
     memo: Dict[Any, Any] = ...
-    def __init__(self, dialect: engine.Dialect, connection: engine.Connection, checkfirst: bool = ...,
+    def __init__(self, dialect: Dialect, connection: Connection, checkfirst: bool = ...,
                  tables: Optional[List[Table]] = ..., **kwargs: Any) -> None: ...
     def visit_metadata(self, metadata: MetaData) -> None: ...
     def visit_index(self, index: str) -> None: ...
@@ -137,7 +137,11 @@ class SchemaDropper(DDLBase):
     def visit_sequence(self, sequence: str, drop_ok: bool = ...) -> None: ...
 
 def sort_tables(tables: Sequence[Table], skip_fn: Optional[Callable[[ForeignKey], bool]] = ...,
-                extra_dependencies: Optional[Sequence[Tuple[Table, Table]]] = ...): ...
+                extra_dependencies: Optional[Sequence[Tuple[Table, Table]]] = ...) -> List[Table]: ...
+
+
+_SortResultType = List[Tuple[Optional[Table], Union[Set[ForeignKeyConstraint], List[ForeignKeyConstraint]]]]
+
 def sort_tables_and_constraints(tables: Sequence[Table],
                                 filter_fn: Optional[Callable[[ForeignKeyConstraint], Optional[bool]]] = ...,
-                                extra_dependencies: Optional[Sequence[Tuple[Table, Table]]] = ...): ...
+                                extra_dependencies: Optional[Sequence[Tuple[Table, Table]]] = ...) -> _SortResultType: ...

sqlalchemy-stubs/sql/dml.pyi

@@ -1,51 +1,65 @@
-from typing import Any, Optional
+from typing import Any, Optional, Union, TypeVar, Iterable, Tuple
 from .base import Executable as Executable, DialectKWArgs as DialectKWArgs
-from .elements import ClauseElement as ClauseElement
-from .selectable import HasPrefixes as HasPrefixes, HasCTE as HasCTE
+from .elements import ClauseElement as ClauseElement, ColumnElement, TextClause
+from .selectable import HasPrefixes as HasPrefixes, HasCTE as HasCTE, TableClause, Selectable
+from .schema import Table, Column
+from .visitors import Visitable
+
+_UB = TypeVar('_UB', bound=UpdateBase)
 
 class UpdateBase(HasCTE, DialectKWArgs, HasPrefixes, Executable, ClauseElement):
     __visit_name__: str = ...
     named_with_column: bool = ...
-    def params(self, *arg, **kw): ...
+    def params(self, *arg: Any, **kw: Any) -> Any: ...
     bind: Any = ...
-    def returning(self, *cols): ...
-    def with_hint(self, text, selectable: Optional[Any] = ..., dialect_name: str = ...): ...
+    def returning(self: _UB, *cols: ColumnElement[Any]) -> _UB: ...
+    def with_hint(self: _UB, text: str, selectable: Optional[Table] = ..., dialect_name: str = ...) -> _UB: ...
+
+_VB = TypeVar('_VB', bound=ValuesBase)
 
 class ValuesBase(UpdateBase):
     __visit_name__: str = ...
     select: Any = ...
-    table: Any = ...
-    def __init__(self, table, values, prefixes) -> None: ...
+    table: Union[TextClause, Selectable] = ...
+    def __init__(self, table: Union[str, Selectable], values: Any, prefixes: Any) -> None: ...
     parameters: Any = ...
-    def values(self, *args, **kwargs): ...
-    def return_defaults(self, *cols): ...
+    def values(self: _VB, *args: Any, **kwargs: Any) -> _VB: ...
+    def return_defaults(self: _VB, *cols: Any) -> _VB: ...
+
+_I = TypeVar('_I', bound=Insert)
 
 class Insert(ValuesBase):
     __visit_name__: str = ...
-    select: Any = ...
+    select: Optional[Selectable] = ...
     include_insert_from_select_defaults: bool = ...
-    inline: Any = ...
-    def __init__(self, table, values: Optional[Any] = ..., inline: bool = ..., bind: Optional[Any] = ...,
+    inline: bool = ...
+    select_names: Optional[Iterable[Union[str, Column[Any]]]] = ...
+    def __init__(self, table: Union[str, Selectable], values: Optional[Any] = ..., inline: bool = ..., bind: Optional[Any] = ...,
                  prefixes: Optional[Any] = ..., returning: Optional[Any] = ..., return_defaults: bool = ...,
-                 **dialect_kw) -> None: ...
-    def get_children(self, **kwargs): ...
-    select_names: Any = ...
-    def from_select(self, names, select, include_defaults: bool = ...): ...
+                 **dialect_kw: Any) -> None: ...
+    def get_children(self, **kwargs: Any) -> Union[Tuple, Tuple[Selectable]]: ...
+    def from_select(self: _I, names: Iterable[Union[str, Column[Any]]], select: Union[str, Selectable],
+                    include_defaults: bool = ...) -> _I: ...
+
+_U = TypeVar('_U', bound=Update)
 
 class Update(ValuesBase):
     __visit_name__: str = ...
-    inline: Any = ...
-    def __init__(self, table, whereclause: Optional[Any] = ..., values: Optional[Any] = ...,
-                 inline: bool = ..., bind: Optional[Any] = ..., prefixes: Optional[Any] = ...,
+    inline: bool = ...
+    def __init__(self, table: Union[str, Selectable], whereclause: Optional[Union[str, bool, Visitable]] = ...,
+                 values: Optional[Any] = ..., inline: bool = ..., bind: Optional[Any] = ..., prefixes: Optional[Any] = ...,
                  returning: Optional[Any] = ..., return_defaults: bool = ..., preserve_parameter_order: bool = ...,
-                 **dialect_kw) -> None: ...
-    def get_children(self, **kwargs): ...
-    def where(self, whereclause): ...
+                 **dialect_kw: Any) -> None: ...
+    def get_children(self, **kwargs: Any) -> Union[Tuple, Tuple[Visitable]]: ...
+    def where(self: _U, whereclause: Union[str, bool, Visitable]) -> _U: ...
+
+_D = TypeVar('_D', bound=Delete)
 
 class Delete(UpdateBase):
     __visit_name__: str = ...
-    table: Any = ...
-    def __init__(self, table, whereclause: Optional[Any] = ..., bind: Optional[Any] = ...,
-                 returning: Optional[Any] = ..., prefixes: Optional[Any] = ..., **dialect_kw) -> None: ...
-    def get_children(self, **kwargs): ...
-    def where(self, whereclause): ...
+    table: Union[TextClause, Selectable] = ...
+    def __init__(self, table: Union[str, Selectable], whereclause: Optional[Union[str, bool, Visitable]] = ...,
+                 bind: Optional[Any] = ..., returning: Optional[Any] = ..., prefixes: Optional[Any] = ...,
+                 **dialect_kw: Any) -> None: ...
+    def get_children(self, **kwargs: Any) -> Union[Tuple, Tuple[Visitable]]: ...
+    def where(self: _D, whereclause: Union[str, bool, Visitable]) -> _D: ...

sqlalchemy-stubs/sql/elements.pyi

@@ -501,7 +501,7 @@ class conv(_truncated_label): ...
 
 class _anonymous_label(_truncated_label): ...
 
-class AnnotatedColumnElement(Annotated, Generic[_T]):
+class AnnotatedColumnElement(Annotated[ColumnElement[_T]], Generic[_T]):
     def __init__(self, element: ColumnElement[_T], values: Any) -> None: ...
     @property
     def name(self) -> str: ...