# ### this file stubs are generated by tools/write_pyi.py - do not edit ###
# ### imports are manually managed
-
+from contextlib import contextmanager
from typing import Any
from typing import Callable
+from typing import Dict
+from typing import Iterator
from typing import List
+from typing import Literal
+from typing import Mapping
from typing import Optional
from typing import Sequence
+from typing import Tuple
from typing import Type
from typing import TYPE_CHECKING
from typing import Union
from sqlalchemy.sql.type_api import TypeEngine
from sqlalchemy.util import immutabledict
+ from .operations.ops import BatchOperations
from .operations.ops import MigrateOperation
from .util.sqla_compat import _literal_bindparam
"""
+@contextmanager
def batch_alter_table(
- table_name,
- schema=None,
- recreate="auto",
- partial_reordering=None,
- copy_from=None,
- table_args=(),
- table_kwargs=immutabledict({}),
- reflect_args=(),
- reflect_kwargs=immutabledict({}),
- naming_convention=None,
-):
+ table_name: str,
+ schema: Optional[str] = None,
+ recreate: Literal["auto", "always", "never"] = "auto",
+ partial_reordering: Optional[tuple] = None,
+ copy_from: Optional["Table"] = None,
+ table_args: Tuple[Any, ...] = (),
+ table_kwargs: Mapping[str, Any] = immutabledict({}),
+ reflect_args: Tuple[Any, ...] = (),
+ reflect_kwargs: Mapping[str, Any] = immutabledict({}),
+ naming_convention: Optional[Dict[str, str]] = None,
+) -> Iterator["BatchOperations"]:
"""Invoke a series of per-table migrations in batch.
Batch mode allows a series of operations specific to a table
import textwrap
from typing import Any
from typing import Callable
+from typing import Dict
from typing import Iterator
from typing import List # noqa
+from typing import Mapping
from typing import Optional
from typing import Sequence # noqa
+from typing import Tuple
from typing import Type # noqa
from typing import TYPE_CHECKING
from typing import Union
NoneType = type(None)
if TYPE_CHECKING:
+ from typing import Literal
+
from sqlalchemy import Table # noqa
from sqlalchemy.engine import Connection
@contextmanager
def batch_alter_table(
self,
- table_name,
- schema=None,
- recreate="auto",
- partial_reordering=None,
- copy_from=None,
- table_args=(),
- table_kwargs=util.immutabledict(),
- reflect_args=(),
- reflect_kwargs=util.immutabledict(),
- naming_convention=None,
- ):
+ table_name: str,
+ schema: Optional[str] = None,
+ recreate: Literal["auto", "always", "never"] = "auto",
+ partial_reordering: Optional[tuple] = None,
+ copy_from: Optional["Table"] = None,
+ table_args: Tuple[Any, ...] = (),
+ table_kwargs: Mapping[str, Any] = util.immutabledict(),
+ reflect_args: Tuple[Any, ...] = (),
+ reflect_kwargs: Mapping[str, Any] = util.immutabledict(),
+ naming_convention: Optional[Dict[str, str]] = None,
+ ) -> Iterator["BatchOperations"]:
"""Invoke a series of per-table migrations in batch.
Batch mode allows a series of operations specific to a table
def importlib_metadata_get(group: str) -> Sequence[EntryPoint]:
ep = importlib_metadata.entry_points()
if hasattr(ep, "select"):
- return ep.select(group=group) # type: ignore
+ return ep.select(group=group)
else:
- return ep.get(group, ()) # type: ignore
+ return ep.get(group, ())
def formatannotation_fwdref(annotation, base_module=None):
"sqlalchemy.sql.functions.",
"sqlalchemy.sql.dml.",
]
+CONTEXT_MANAGERS = {"op": ["batch_alter_table"]}
def generate_pyi_for_proxy(
source_path: Path,
destination_path: Path,
ignore_output: bool,
- ignore_items: set,
+ file_key: str,
):
+ ignore_items = IGNORE_ITEMS.get(file_key, set())
+ context_managers = CONTEXT_MANAGERS.get(file_key, [])
if sys.version_info < (3, 9):
raise RuntimeError("This script must be run with Python 3.9 or higher")
continue
meth = getattr(cls, name, None)
if callable(meth):
- _generate_stub_for_meth(cls, name, printer, env)
+ _generate_stub_for_meth(
+ cls, name, printer, env, name in context_managers
+ )
else:
_generate_stub_for_attr(cls, name, printer, env)
printer.writeline(f"{name}: {type_}")
-def _generate_stub_for_meth(cls, name, printer, env):
+def _generate_stub_for_meth(cls, name, printer, env, is_context_manager):
fn = getattr(cls, name)
while hasattr(fn, "__wrapped__"):
formatannotation=_formatannotation,
formatreturns=lambda val: f"-> {_formatannotation(val)}",
)
-
+ contextmanager = "@contextmanager" if is_context_manager else ""
func_text = textwrap.dedent(
- """\
- def %(name)s%(argspec)s:
- '''%(doc)s'''
+ f"""
+ {contextmanager}
+ def {name}{argspec}:
+ '''{fn.__doc__}'''
"""
- % {
- "name": name,
- "argspec": argspec,
- "doc": fn.__doc__,
- }
)
-
printer.write_indented_block(func_text)
def run_file(
- source_path: Path, cls_to_generate: type, stdout: bool, ignore_items: set
+ source_path: Path, cls_to_generate: type, stdout: bool, file_key: str
):
progname = Path(sys.argv[0]).as_posix()
if not stdout:
source_path=source_path,
destination_path=source_path,
ignore_output=False,
- ignore_items=ignore_items,
+ file_key=file_key,
)
else:
with NamedTemporaryFile(delete=False, suffix=".pyi") as f:
source_path=source_path,
destination_path=f_path,
ignore_output=True,
- ignore_items=ignore_items,
+ file_key=file_key,
)
sys.stdout.write(f_path.read_text())
f_path.unlink()
def main(args):
location = Path(__file__).parent.parent / "alembic"
if args.file in {"all", "op"}:
- run_file(
- location / "op.pyi", Operations, args.stdout, IGNORE_ITEMS["op"]
- )
+ run_file(location / "op.pyi", Operations, args.stdout, "op")
if args.file in {"all", "context"}:
run_file(
location / "context.pyi",
EnvironmentContext,
args.stdout,
- IGNORE_ITEMS["context"],
+ "context",
)