Add Operations and BatchOperations stub methods

Updated stub generator script to also add stubs method definitions
for the :class:`.Operations` class and the :class:`.BatchOperations`
class obtained from :meth:`.Operations.batch_alter_table`.

Repaired the return signatures for :class:`.Operations` that mostly
return ``None``, and were erroneously referring to ``Optional[Table]``
in many cases.

Fixes: #1093
Change-Id: I98d38dd5a1e719b4dbbc1003746ec28f26c27808

alembic/context.pyi

@@ -3,10 +3,14 @@
 from __future__ import annotations
 
 from typing import Any
+from typing import Callable
+from typing import Collection
 from typing import ContextManager
 from typing import Dict
 from typing import List
 from typing import Literal
+from typing import Mapping
+from typing import MutableMapping
 from typing import Optional
 from typing import overload
 from typing import TextIO
@@ -19,16 +23,16 @@ if TYPE_CHECKING:
     from sqlalchemy.engine.url import URL
     from sqlalchemy.sql.elements import ClauseElement
     from sqlalchemy.sql.schema import MetaData
+    from sqlalchemy.sql.schema import SchemaItem
 
+    from .autogenerate.api import AutogenContext
     from .config import Config
-    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 .operations.ops import MigrateOperation
     from .runtime.migration import _ProxyTransaction
     from .runtime.migration import MigrationContext
+    from .runtime.migration import MigrationInfo
     from .script import ScriptDirectory
+
 ### end imports ###
 
 def begin_transaction() -> Union[_ProxyTransaction, ContextManager[None]]:
@@ -79,7 +83,7 @@ config: Config
 
 def configure(
     connection: Optional[Connection] = None,
-    url: Optional[Union[str, URL]] = None,
+    url: Union[str, URL, None] = None,
     dialect_name: Optional[str] = None,
     dialect_opts: Optional[Dict[str, Any]] = None,
     transactional_ddl: Optional[bool] = None,
@@ -90,20 +94,77 @@ def configure(
     template_args: Optional[Dict[str, Any]] = None,
     render_as_batch: bool = False,
     target_metadata: Optional[MetaData] = None,
-    include_name: Optional[IncludeNameFn] = None,
-    include_object: Optional[IncludeObjectFn] = None,
+    include_name: Optional[
+        Callable[
+            [
+                Optional[str],
+                Literal[
+                    "schema",
+                    "table",
+                    "column",
+                    "index",
+                    "unique_constraint",
+                    "foreign_key_constraint",
+                ],
+                MutableMapping[
+                    Literal[
+                        "schema_name",
+                        "table_name",
+                        "schema_qualified_table_name",
+                    ],
+                    Optional[str],
+                ],
+            ],
+            bool,
+        ]
+    ] = None,
+    include_object: Optional[
+        Callable[
+            [
+                SchemaItem,
+                Optional[str],
+                Literal[
+                    "schema",
+                    "table",
+                    "column",
+                    "index",
+                    "unique_constraint",
+                    "foreign_key_constraint",
+                ],
+                bool,
+                Optional[SchemaItem],
+            ],
+            bool,
+        ]
+    ] = None,
     include_schemas: bool = False,
-    process_revision_directives: Optional[ProcessRevisionDirectiveFn] = None,
+    process_revision_directives: Optional[
+        Callable[
+            [MigrationContext, Tuple[str, str], List[MigrateOperation]], None
+        ]
+    ] = None,
     compare_type: bool = False,
     compare_server_default: bool = False,
-    render_item: Optional[RenderItemFn] = None,
+    render_item: Optional[
+        Callable[[str, Any, AutogenContext], Union[str, Literal[False]]]
+    ] = 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[OnVersionApplyFn] = None,
+    on_version_apply: Optional[
+        Callable[
+            [
+                MigrationContext,
+                MigrationInfo,
+                Collection[Any],
+                Mapping[str, Any],
+            ],
+            None,
+        ]
+    ] = None,
     **kw: Any,
 ) -> None:
     """Configure a :class:`.MigrationContext` within this

alembic/ddl/postgresql.py

@@ -593,7 +593,11 @@ def create_exclude_constraint(
 
     @classmethod
     def batch_create_exclude_constraint(
-        cls, operations, constraint_name, *elements, **kw
+        cls,
+        operations: BatchOperations,
+        constraint_name: str,
+        *elements: Any,
+        **kw: Any,
     ):
         """Issue a "create exclude constraint" instruction using the
         current batch migration context.

alembic/op.pyi

@@ -42,7 +42,7 @@ if TYPE_CHECKING:
 
 def add_column(
     table_name: str, column: Column, schema: Optional[str] = None
-) -> Optional[Table]:
+) -> None:
     """Issue an "add column" instruction using the current
     migration context.
 
@@ -60,19 +60,19 @@ def add_column(
 
     .. note::
 
-        With the exception of NOT NULL constraints or single-column FOREIGN KEY
-        constraints, other kinds of constraints such as PRIMARY KEY, UNIQUE or
-        CHECK constraints **cannot** be generated using this method; for these
-        constraints, refer to operations such as
+        With the exception of NOT NULL constraints or single-column FOREIGN
+        KEY constraints, other kinds of constraints such as PRIMARY KEY,
+        UNIQUE or CHECK constraints **cannot** be generated using this
+        method; for these constraints, refer to operations such as
         :meth:`.Operations.create_primary_key` and
         :meth:`.Operations.create_check_constraint`. In particular, the
         following :class:`~sqlalchemy.schema.Column` parameters are
         **ignored**:
 
         * :paramref:`~sqlalchemy.schema.Column.primary_key` - SQL databases
-          typically do not support an ALTER operation that can add individual
-          columns one at a time to an existing primary key constraint,
-          therefore it's less ambiguous to use the
+          typically do not support an ALTER operation that can add
+          individual columns one at a time to an existing primary key
+          constraint, therefore it's less ambiguous to use the
           :meth:`.Operations.create_primary_key` method, which assumes no
           existing primary key constraint is present.
         * :paramref:`~sqlalchemy.schema.Column.unique` - use the
@@ -137,7 +137,7 @@ def alter_column(
     existing_comment: Optional[str] = None,
     schema: Optional[str] = None,
     **kw: Any,
-) -> Optional[Table]:
+) -> None:
     r"""Issue an "alter column" instruction using the
     current migration context.
 
@@ -483,10 +483,10 @@ def bulk_insert(
 def create_check_constraint(
     constraint_name: Optional[str],
     table_name: str,
-    condition: Union[str, BinaryExpression],
+    condition: Union[str, BinaryExpression, TextClause],
     schema: Optional[str] = None,
     **kw: Any,
-) -> Optional[Table]:
+) -> None:
     """Issue a "create check constraint" instruction using the
     current migration context.
 
@@ -580,7 +580,7 @@ def create_foreign_key(
     source_schema: Optional[str] = None,
     referent_schema: Optional[str] = None,
     **dialect_kw: Any,
-) -> Optional[Table]:
+) -> None:
     """Issue a "create foreign key" instruction using the
     current migration context.
 
@@ -638,7 +638,7 @@ def create_index(
     schema: Optional[str] = None,
     unique: bool = False,
     **kw: Any,
-) -> Optional[Table]:
+) -> None:
     r"""Issue a "create index" instruction using the current
     migration context.
 
@@ -688,7 +688,7 @@ def create_primary_key(
     table_name: str,
     columns: List[str],
     schema: Optional[str] = None,
-) -> Optional[Table]:
+) -> None:
     """Issue a "create primary key" instruction using the current
     migration context.
 
@@ -724,9 +724,7 @@ def create_primary_key(
 
     """
 
-def create_table(
-    table_name: str, *columns: SchemaItem, **kw: Any
-) -> Optional[Table]:
+def create_table(table_name: str, *columns: SchemaItem, **kw: Any) -> Table:
     r"""Issue a "create table" instruction using the current migration
     context.
 
@@ -810,7 +808,7 @@ def create_table_comment(
     comment: Optional[str],
     existing_comment: Optional[str] = None,
     schema: Optional[str] = None,
-) -> Optional[Table]:
+) -> None:
     """Emit a COMMENT ON operation to set the comment for a table.
 
     .. versionadded:: 1.0.6
@@ -878,7 +876,7 @@ def create_unique_constraint(
 
 def drop_column(
     table_name: str, column_name: str, schema: Optional[str] = None, **kw: Any
-) -> Optional[Table]:
+) -> None:
     """Issue a "drop column" instruction using the current
     migration context.
 
@@ -921,7 +919,7 @@ def drop_constraint(
     table_name: str,
     type_: Optional[str] = None,
     schema: Optional[str] = None,
-) -> Optional[Table]:
+) -> None:
     r"""Drop a constraint of the given name, typically via DROP CONSTRAINT.
 
     :param constraint_name: name of the constraint.
@@ -940,7 +938,7 @@ def drop_index(
     table_name: Optional[str] = None,
     schema: Optional[str] = None,
     **kw: Any,
-) -> Optional[Table]:
+) -> None:
     r"""Issue a "drop index" instruction using the current
     migration context.
 
@@ -988,7 +986,7 @@ def drop_table_comment(
     table_name: str,
     existing_comment: Optional[str] = None,
     schema: Optional[str] = None,
-) -> Optional[Table]:
+) -> None:
     """Issue a "drop table comment" operation to
     remove an existing comment set on a table.
 
@@ -1009,7 +1007,7 @@ def drop_table_comment(
 def execute(
     sqltext: Union[str, TextClause, Update],
     execution_options: Optional[dict[str, Any]] = None,
-) -> Optional[Table]:
+) -> None:
     r"""Execute the given SQL using the current migration context.
 
     The given SQL can be a plain string, e.g.::
@@ -1177,8 +1175,8 @@ def inline_literal(
     advanced types like dates may not be supported directly
     by SQLAlchemy.
 
-    See :meth:`.execute` for an example usage of
-    :meth:`.inline_literal`.
+    See :meth:`.Operations.execute` for an example usage of
+    :meth:`.Operations.inline_literal`.
 
     The environment can also be configured to attempt to render
     "literal" values inline automatically, for those simple types
@@ -1229,7 +1227,7 @@ def register_operation(
 
 def rename_table(
     old_table_name: str, new_table_name: str, schema: Optional[str] = None
-) -> Optional[Table]:
+) -> None:
     """Emit an ALTER TABLE to rename a table.
 
     :param old_table_name: old name.

alembic/operations/__init__.py

@@ -1,7 +1,13 @@
 from . import toimpl
+from .base import AbstractOperations
 from .base import BatchOperations
 from .base import Operations
 from .ops import MigrateOperation
 
 
-__all__ = ["Operations", "BatchOperations", "MigrateOperation"]
+__all__ = [
+    "AbstractOperations",
+    "Operations",
+    "BatchOperations",
+    "MigrateOperation",
+]