Merge "Improve typing." into main

alembic/autogenerate/api.py

@@ -9,7 +9,6 @@
 from typing import Set
 from typing import Tuple
 from typing import TYPE_CHECKING
-from typing import Union
 
 from sqlalchemy import inspect
 
@@ -25,19 +24,18 @@
     from sqlalchemy.engine import Connection
     from sqlalchemy.engine import Dialect
     from sqlalchemy.engine import Inspector
-    from sqlalchemy.sql.schema import Column
-    from sqlalchemy.sql.schema import ForeignKeyConstraint
-    from sqlalchemy.sql.schema import Index
     from sqlalchemy.sql.schema import MetaData
-    from sqlalchemy.sql.schema import Table
-    from sqlalchemy.sql.schema import UniqueConstraint
+    from sqlalchemy.sql.schema import SchemaItem
 
-    from alembic.config import Config
-    from alembic.operations.ops import MigrationScript
-    from alembic.operations.ops import UpgradeOps
-    from alembic.runtime.migration import MigrationContext
-    from alembic.script.base import Script
-    from alembic.script.base import ScriptDirectory
+    from ..config import Config
+    from ..operations.ops import MigrationScript
+    from ..operations.ops import UpgradeOps
+    from ..runtime.environment import NameFilterParentNames
+    from ..runtime.environment import NameFilterType
+    from ..runtime.environment import RenderItemFn
+    from ..runtime.migration import MigrationContext
+    from ..script.base import Script
+    from ..script.base import ScriptDirectory
 
 
 def compare_metadata(context: MigrationContext, metadata: MetaData) -> Any:
@@ -172,7 +170,7 @@ def render_python_code(
     alembic_module_prefix: str = "op.",
     render_as_batch: bool = False,
     imports: Tuple[str, ...] = (),
-    render_item: None = None,
+    render_item: Optional[RenderItemFn] = None,
     migration_context: Optional[MigrationContext] = None,
 ) -> str:
     """Render Python code given an :class:`.UpgradeOps` or
@@ -359,8 +357,8 @@ def _within_batch(self) -> Iterator[None]:
     def run_name_filters(
         self,
         name: Optional[str],
-        type_: str,
-        parent_names: Dict[str, Optional[str]],
+        type_: NameFilterType,
+        parent_names: NameFilterParentNames,
     ) -> bool:
         """Run the context's name filters and return True if the targets
         should be part of the autogenerate operation.
@@ -396,17 +394,11 @@ def run_name_filters(
 
     def run_object_filters(
         self,
-        object_: Union[
-            Table,
-            Index,
-            Column,
-            UniqueConstraint,
-            ForeignKeyConstraint,
-        ],
+        object_: SchemaItem,
         name: Optional[str],
-        type_: str,
+        type_: NameFilterType,
         reflected: bool,
-        compare_to: Optional[Union[Table, Index, Column, UniqueConstraint]],
+        compare_to: Optional[SchemaItem],
     ) -> bool:
         """Run the context's object filters and return True if the targets
         should be part of the autogenerate operation.

alembic/autogenerate/compare.py

@@ -212,7 +212,7 @@ def _compare_tables(
                 (inspector),
                 # fmt: on
             )
-            sqla_compat._reflect_table(inspector, t, None)
+            sqla_compat._reflect_table(inspector, t)
         if autogen_context.run_object_filters(t, tname, "table", True, None):
 
             modify_table_ops = ops.ModifyTableOps(tname, [], schema=s)
@@ -243,7 +243,7 @@ def _compare_tables(
                 _compat_autogen_column_reflect(inspector),
                 # fmt: on
             )
-            sqla_compat._reflect_table(inspector, t, None)
+            sqla_compat._reflect_table(inspector, t)
         conn_column_info[(s, tname)] = t
 
     for s, tname in sorted(existing_tables, key=lambda x: (x[0] or "", x[1])):

alembic/context.pyi

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 from typing import Any
-from typing import Callable
 from typing import ContextManager
 from typing import Dict
 from typing import List
@@ -22,7 +21,11 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.schema import MetaData
 
     from .config import Config
-    from .operations import MigrateOperation
+    from .runtime.environment import IncludeNameFn
+    from .runtime.environment import IncludeObjectFn
+    from .runtime.environment import OnVersionApplyFn
+    from .runtime.environment import ProcessRevisionDirectiveFn
+    from .runtime.environment import RenderItemFn
     from .runtime.migration import _ProxyTransaction
     from .runtime.migration import MigrationContext
     from .script import ScriptDirectory
@@ -76,7 +79,7 @@ config: Config
 
 def configure(
     connection: Optional[Connection] = None,
-    url: Union[str, URL, None] = None,
+    url: Optional[Union[str, URL]] = None,
     dialect_name: Optional[str] = None,
     dialect_opts: Optional[Dict[str, Any]] = None,
     transactional_ddl: Optional[bool] = None,
@@ -87,24 +90,20 @@ def configure(
     template_args: Optional[Dict[str, Any]] = None,
     render_as_batch: bool = False,
     target_metadata: Optional[MetaData] = None,
-    include_name: Optional[Callable[..., bool]] = None,
-    include_object: Optional[Callable[..., bool]] = None,
+    include_name: Optional[IncludeNameFn] = None,
+    include_object: Optional[IncludeObjectFn] = None,
     include_schemas: bool = False,
-    process_revision_directives: Optional[
-        Callable[
-            [MigrationContext, Tuple[str, str], List[MigrateOperation]], None
-        ]
-    ] = None,
+    process_revision_directives: Optional[ProcessRevisionDirectiveFn] = None,
     compare_type: bool = False,
     compare_server_default: bool = False,
-    render_item: Optional[Callable[..., bool]] = None,
+    render_item: Optional[RenderItemFn] = None,
     literal_binds: bool = False,
     upgrade_token: str = "upgrades",
     downgrade_token: str = "downgrades",
     alembic_module_prefix: str = "op.",
     sqlalchemy_module_prefix: str = "sa.",
     user_module_prefix: Optional[str] = None,
-    on_version_apply: Optional[Callable[..., None]] = None,
+    on_version_apply: Optional[OnVersionApplyFn] = None,
     **kw: Any,
 ) -> None:
     """Configure a :class:`.MigrationContext` within this
@@ -308,7 +307,8 @@ def configure(
        ``"unique_constraint"``, or ``"foreign_key_constraint"``
      * ``parent_names``: a dictionary of "parent" object names, that are
        relative to the name being given.  Keys in this dictionary may
-       include:  ``"schema_name"``, ``"table_name"``.
+       include:  ``"schema_name"``, ``"table_name"`` or
+       ``"schema_qualified_table_name"``.
 
      E.g.::
 

alembic/ddl/impl.py

@@ -155,9 +155,9 @@ def bind(self) -> Optional[Connection]:
     def _exec(
         self,
         construct: Union[ClauseElement, str],
-        execution_options: Optional[dict] = None,
+        execution_options: Optional[dict[str, Any]] = None,
         multiparams: Sequence[dict] = (),
-        params: Dict[str, int] = util.immutabledict(),
+        params: Dict[str, Any] = util.immutabledict(),
     ) -> Optional[CursorResult]:
         if isinstance(construct, str):
             construct = text(construct)
@@ -197,7 +197,7 @@ def _exec(
     def execute(
         self,
         sql: Union[ClauseElement, str],
-        execution_options: None = None,
+        execution_options: Optional[dict[str, Any]] = None,
     ) -> None:
         self._exec(sql, execution_options)
 

alembic/op.pyi

@@ -951,7 +951,8 @@ def drop_table_comment(
     """
 
 def execute(
-    sqltext: Union[str, TextClause, Update], execution_options: None = None
+    sqltext: Union[str, TextClause, Update],
+    execution_options: Optional[dict[str, Any]] = None,
 ) -> Optional[Table]:
     r"""Execute the given SQL using the current migration context.
 
@@ -1101,7 +1102,7 @@ def implementation_for(op_cls: Any) -> Callable[..., Any]:
     """
 
 def inline_literal(
-    value: Union[str, int], type_: None = None
+    value: Union[str, int], type_: Optional[TypeEngine] = None
 ) -> _literal_bindparam:
     r"""Produce an 'inline literal' expression, suitable for
     using in an INSERT, UPDATE, or DELETE statement.

alembic/operations/base.py

@@ -33,8 +33,9 @@
 if TYPE_CHECKING:
     from typing import Literal
 
-    from sqlalchemy import Table  # noqa
+    from sqlalchemy import Table
     from sqlalchemy.engine import Connection
+    from sqlalchemy.types import TypeEngine
 
     from .batch import BatchOperationsImpl
     from .ops import MigrateOperation
@@ -439,7 +440,7 @@ def f(self, name: str) -> conv:
         return conv(name)
 
     def inline_literal(
-        self, value: Union[str, int], type_: None = None
+        self, value: Union[str, int], type_: Optional[TypeEngine[Any]] = None
     ) -> _literal_bindparam:
         r"""Produce an 'inline literal' expression, suitable for
         using in an INSERT, UPDATE, or DELETE statement.

alembic/operations/batch.py

@@ -487,7 +487,7 @@ def alter_column(
         server_default: Optional[Union[Function[Any], str, bool]] = False,
         name: Optional[str] = None,
         type_: Optional[TypeEngine] = None,
-        autoincrement: None = None,
+        autoincrement: Optional[Union[bool, Literal["auto"]]] = None,
         comment: Union[str, Literal[False]] = False,
         **kw,
     ) -> None:

alembic/operations/ops.py

@@ -673,11 +673,11 @@ def batch_create_foreign_key(
         local_cols: List[str],
         remote_cols: List[str],
         referent_schema: Optional[str] = None,
-        onupdate: None = None,
-        ondelete: None = None,
-        deferrable: None = None,
-        initially: None = None,
-        match: None = None,
+        onupdate: Optional[str] = None,
+        ondelete: Optional[str] = None,
+        deferrable: Optional[bool] = None,
+        initially: Optional[str] = None,
+        match: Optional[str] = None,
         **dialect_kw: Any,
     ) -> None:
         """Issue a "create foreign key" instruction using the
@@ -1890,10 +1890,10 @@ def batch_alter_column(
         type_: Optional[Union[TypeEngine, Type[TypeEngine]]] = None,
         existing_type: Optional[Union[TypeEngine, Type[TypeEngine]]] = None,
         existing_server_default: bool = False,
-        existing_nullable: None = None,
-        existing_comment: None = None,
-        insert_before: None = None,
-        insert_after: None = None,
+        existing_nullable: Optional[bool] = None,
+        existing_comment: Optional[str] = None,
+        insert_before: Optional[str] = None,
+        insert_after: Optional[str] = None,
         **kw: Any,
     ) -> Optional[Table]:
         """Issue an "alter column" instruction using the current
@@ -1935,6 +1935,8 @@ def batch_alter_column(
             modify_server_default=server_default,
             modify_nullable=nullable,
             modify_comment=comment,
+            insert_before=insert_before,
+            insert_after=insert_after,
             **kw,
         )
 
@@ -2314,7 +2316,7 @@ class ExecuteSQLOp(MigrateOperation):
     def __init__(
         self,
         sqltext: Union[Update, str, Insert, TextClause],
-        execution_options: None = None,
+        execution_options: Optional[dict[str, Any]] = None,
     ) -> None:
         self.sqltext = sqltext
         self.execution_options = execution_options
@@ -2324,7 +2326,7 @@ def execute(
         cls,
         operations: Operations,
         sqltext: Union[str, TextClause, Update],
-        execution_options: None = None,
+        execution_options: Optional[dict[str, Any]] = None,
     ) -> Optional[Table]:
         r"""Execute the given SQL using the current migration context.