Improvements to type to support alembic (#131)

* Improvements to type to support alembic

* Move improvements.

added class var type - instance instance descriptor thingy

* more updates.

* final updates

Author: CaselIT

sqlalchemy-stubs/_typing.pyi

@@ -0,0 +1,17 @@
+from typing import Any
+from typing import Generic
+from typing import overload
+from typing import Type
+from typing import TypeVar
+
+_T = TypeVar("_T")
+
+class _TypeToInstance(Generic[_T]):
+    @overload
+    def __get__(self, instance: None, owner: Any) -> Type[_T]: ...
+    @overload
+    def __get__(self, instance: object, owner: Any) -> _T: ...
+    @overload
+    def __set__(self, instance: None, value: Type[_T]) -> None: ...
+    @overload
+    def __set__(self, instance: object, value: _T) -> None: ...

sqlalchemy-stubs/engine/base.pyi

@@ -69,7 +69,7 @@ class Connection(Connectable):
     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) -> Optional[Transaction]: ...
+    def begin(self) -> Transaction: ...
     def begin_nested(self) -> NestedTransaction: ...
     def begin_twophase(
         self, xid: Optional[str] = ...

sqlalchemy-stubs/engine/default.pyi

@@ -16,6 +16,7 @@ from .cursor import CursorFetchStrategy
 from .url import URL
 from .. import pool
 from .. import util
+from .._typing import _TypeToInstance
 from ..sql import compiler
 from ..sql.type_api import TypeEngine
 from ..util import langhelpers
@@ -30,14 +31,9 @@ CACHING_DISABLED: langhelpers._symbol
 NO_CACHE_KEY: langhelpers._symbol
 
 class DefaultDialect(interfaces.Dialect):
-    statement_compiler: Type[compiler.SQLCompiler] = ...
-    ddl_compiler: Type[compiler.DDLCompiler] = ...
-    type_compiler: Type[compiler.GenericTypeCompiler] = ...
-    preparer: Type[compiler.IdentifierPreparer] = ...
+    type_compiler: _TypeToInstance[compiler.GenericTypeCompiler] = ...
     execution_ctx_cls: Type[DefaultExecutionContext] = ...
     supports_alter: bool = ...
-    supports_comments: bool = ...
-    inline_comments: bool = ...
     use_setinputsizes: bool = ...
     default_sequence_base: int = ...
     execute_sequence_format: Union[
@@ -96,7 +92,6 @@ class DefaultDialect(interfaces.Dialect):
     positional: bool = ...
     dbapi: Any = ...
     paramstyle: str = ...
-    identifier_preparer: compiler.IdentifierPreparer = ...
     case_sensitive: bool = ...
     label_length: Optional[int] = ...
     compiler_linting: int = ...

sqlalchemy-stubs/engine/interfaces.pyi

@@ -16,7 +16,9 @@ from .base import Engine
 from .cursor import CursorFetchStrategy
 from .cursor import CursorResult
 from .url import URL
+from .._typing import _TypeToInstance
 from ..exc import StatementError
+from ..sql import compiler
 from ..sql.compiler import Compiled as Compiled
 from ..sql.compiler import IdentifierPreparer
 from ..sql.compiler import TypeCompiler as TypeCompiler
@@ -92,15 +94,18 @@ class Dialect:
     positional: bool = ...
     paramstyle: str = ...
     encoding: str = ...
-    statement_compiler: Type[Compiled] = ...
-    ddl_compiler: Type[Compiled] = ...
+    statement_compiler: Type[compiler.SQLCompiler] = ...
+    ddl_compiler: Type[compiler.DDLCompiler] = ...
+    type_compiler: _TypeToInstance[compiler.TypeCompiler] = ...
+    preparer: Type[compiler.IdentifierPreparer] = ...
+    identifier_preparer: compiler.IdentifierPreparer = ...
     server_version_info: Optional[Tuple[Any, ...]] = ...
     default_schema_name: Optional[str] = ...
     execution_ctx_cls: Type[ExecutionContext] = ...
     execute_sequence_format: Union[
         Type[Tuple[Any, ...]], Type[List[Any]]
     ] = ...
-    preparer: Type[IdentifierPreparer] = ...
+    preparer: IdentifierPreparer = ...
     supports_alter: bool = ...
     max_identifier_length: int = ...
     supports_sane_rowcount: bool = ...
@@ -114,6 +119,9 @@ class Dialect:
     supports_native_boolean: bool = ...
     dbapi_exception_translation_map: Dict[str, str] = ...
     supports_statement_cache: bool = ...
+    supports_comments: bool = ...
+    inline_comments: bool = ...
+    def __init__(self, *args: Any, **kw: Any) -> None: ...
     def create_connect_args(
         self, url: URL
     ) -> Tuple[Sequence[Any], Mapping[str, Any]]: ...

sqlalchemy-stubs/engine/reflection.pyi

@@ -7,15 +7,22 @@ from typing import Set
 from typing import Tuple
 from typing import Type
 from typing import TypeVar
+from typing import Union
 
 from .base import Connectable
+from .base import Connection
+from .base import Engine
+from .interfaces import Dialect
 from ..sql.schema import Table
 
 _TInspector = TypeVar("_TInspector", bound=Inspector)
 
 def cache(fn: Any, self: Any, con: Any, *args: Any, **kw: Any) -> Any: ...
 
 class Inspector:
+    bind: Union[Engine, Connection] = ...
+    engine: Engine = ...
+    dialect: Dialect = ...
     def __init__(self, bind: Connectable): ...
     @classmethod
     def from_engine(
@@ -68,15 +75,15 @@ class Inspector:
     def reflecttable(
         self,
         table: Table,
-        include_columns: Collection[str],
+        include_columns: Optional[Collection[str]],
         exclude_columns: Collection[str] = ...,
         resolve_fks: bool = ...,
         _extend_on: Set[Table] = ...,
     ) -> None: ...
     def reflect_table(
         self,
         table: Table,
-        include_columns: Collection[str],
+        include_columns: Optional[Collection[str]],
         exclude_columns: Collection[str] = ...,
         resolve_fks: bool = ...,
         _extend_on: Set[Table] = ...,

sqlalchemy-stubs/orm/__init__.pyi

@@ -3,8 +3,8 @@ from typing import Mapping
 from typing import Optional
 from typing import Tuple
 
-from . import exc
-from . import mapper as mapperlib
+from . import exc as exc
+from . import mapper as mapperlib  # noqa
 from . import strategy_options
 from .attributes import AttributeEvent as AttributeEvent
 from .attributes import InstrumentedAttribute as InstrumentedAttribute

sqlalchemy-stubs/schema.pyi

@@ -21,6 +21,7 @@ from .sql.ddl import SetTableComment as SetTableComment
 from .sql.ddl import sort_tables as sort_tables
 from .sql.ddl import sort_tables_and_constraints as sort_tables_and_constraints
 from .sql.naming import conv as conv
+from .sql.schema import _get_table_key as _get_table_key
 from .sql.schema import BLANK_SCHEMA as BLANK_SCHEMA
 from .sql.schema import CheckConstraint as CheckConstraint
 from .sql.schema import Column as Column

sqlalchemy-stubs/sql/compiler.pyi

@@ -94,7 +94,7 @@ class Compiled:
         self, element: Any, err: Any
     ) -> _VisitResult: ...
     @property
-    def sql_compiler(self) -> Compiled: ...
+    def sql_compiler(self) -> SQLCompiler: ...
     def process(self, obj: Any, **kwargs: Any) -> _VisitResult: ...
     def construct_params(
         self,

sqlalchemy-stubs/sql/elements.pyi

@@ -714,6 +714,9 @@ class AnnotatedColumnElement(Annotated):
     def anon_label(self) -> Any: ...
 
 class _truncated_label(quoted_name):
+    def __new__(
+        cls: Type[_QN], value: Any, quote: Optional[Any] = ...
+    ) -> _QN: ...
     def __reduce__(self) -> Any: ...
     def apply_map(self: _TL, map_: Any) -> _TL: ...
 

sqlalchemy-stubs/sql/schema.pyi

@@ -8,6 +8,7 @@ from typing import List
 from typing import Mapping
 from typing import Optional
 from typing import overload
+from typing import Sequence
 from typing import Set
 from typing import Type
 from typing import TypeVar
@@ -22,11 +23,13 @@ from . import type_api
 from . import visitors
 from .base import DedupeColumnCollection
 from .base import DialectKWArgs
+from .base import ImmutableColumnCollection
 from .base import SchemaEventTarget
 from .elements import ClauseElement
 from .elements import ColumnClause
 from .elements import ColumnElement
 from .elements import TextClause
+from .events import DDLEvents
 from .selectable import TableClause
 from .. import util
 from ..engine import Connection
@@ -54,6 +57,7 @@ _ID = TypeVar("_ID", bound=Identity)
 class SchemaItem(SchemaEventTarget, visitors.Visitable):
     __visit_name__: str = ...
     create_drop_stringify_dialect: str = ...
+    dispatch: DDLEvents
     @util.memoized_property
     def info(self) -> Dict[Any, Any]: ...
 
@@ -67,7 +71,8 @@ class Table(DialectKWArgs, SchemaItem, TableClause):
     primary_key: PrimaryKeyConstraint = ...
     fullname: str = ...
     implicit_returning: bool = ...
-    comment: Optional[Any] = ...
+    comment: Optional[str] = ...
+    _prefixes: Sequence[str] = ...
     def __new__(cls: Type[_TAB], *args: Any, **kw: Any) -> _TAB: ...
     def __init__(self, *args: Any, **kw: Any) -> None: ...
     @property
@@ -118,9 +123,18 @@ class Table(DialectKWArgs, SchemaItem, TableClause):
         ] = ...,
         name: Optional[util.text_type] = ...,
     ) -> Table: ...
+    @util.memoized_property
+    def columns(self) -> ImmutableColumnCollection[Column[Any]]: ...
+    @util.memoized_property
+    def c(self) -> ImmutableColumnCollection[Column[Any]]: ...
+    @property
+    def _autoincrement_column(self) -> Optional[Column[Any]]: ...
+    @_autoincrement_column.setter
+    def _autoincrement_column(self, value: Optional[Column[Any]]) -> None: ...
 
 class Column(DialectKWArgs, SchemaItem, ColumnClause[_TE]):
     __visit_name__: str = ...
+    name: str = ...
     inherit_cache: bool = ...
     key: Optional[str] = ...
     primary_key: bool = ...
@@ -438,6 +452,7 @@ class Constraint(DialectKWArgs, SchemaItem):
     name: Optional[str] = ...
     deferrable: Optional[bool] = ...
     initially: Optional[str] = ...
+    parent: Optional[Table] = ...
     def __init__(
         self,
         name: Optional[str] = ...,
@@ -502,7 +517,7 @@ class ForeignKeyConstraint(ColumnCollectionConstraint):
     elements: List[ForeignKey] = ...
     def __init__(
         self,
-        columns: Iterable[str],
+        columns: Iterable[Union[str, Column[Any]]],
         refcolumns: Iterable[Union[str, Column[Any]]],
         name: Optional[str] = ...,
         onupdate: Optional[str] = ...,
@@ -651,3 +666,5 @@ class Identity(IdentityOptions, FetchedValue, SchemaItem):
         order: Optional[bool] = ...,
     ) -> None: ...
     def copy(self: _ID, **kw: Any) -> _ID: ...
+
+def _get_table_key(name: str, schema: Optional[str]) -> str: ...