diff options
| -rw-r--r-- | alembic/config.py | 4 | ||||
| -rw-r--r-- | alembic/context.pyi | 726 | ||||
| -rw-r--r-- | alembic/op.pyi | 1160 | ||||
| -rw-r--r-- | alembic/operations/base.py | 15 | ||||
| -rw-r--r-- | alembic/py.typed | 0 | ||||
| -rw-r--r-- | alembic/script/write_hooks.py | 7 | ||||
| -rw-r--r-- | alembic/testing/__init__.py | 1 | ||||
| -rw-r--r-- | alembic/util/compat.py | 114 | ||||
| -rw-r--r-- | alembic/util/langhelpers.py | 5 | ||||
| -rw-r--r-- | alembic/util/sqla_compat.py | 4 | ||||
| -rw-r--r-- | docs/build/unreleased/py3_typing.rst | 12 | ||||
| -rw-r--r-- | setup.cfg | 3 | ||||
| -rw-r--r-- | tests/test_stubs.py | 52 | ||||
| -rw-r--r-- | tools/write_pyi.py | 195 | ||||
| -rw-r--r-- | tox.ini | 2 |
15 files changed, 2174 insertions, 126 deletions
diff --git a/alembic/config.py b/alembic/config.py index dbcd106..273acbb 100644 --- a/alembic/config.py +++ b/alembic/config.py @@ -520,8 +520,8 @@ class CommandLine: and fn.__module__ == "alembic.command" ): - spec = compat.inspect_getargspec(fn) - if spec[3]: + spec = compat.inspect_getfullargspec(fn) + if spec[3] is not None: positional = spec[0][1 : -len(spec[3])] kwarg = spec[0][-len(spec[3]) :] else: diff --git a/alembic/context.pyi b/alembic/context.pyi new file mode 100644 index 0000000..8a33d52 --- /dev/null +++ b/alembic/context.pyi @@ -0,0 +1,726 @@ +# ### this file stubs are generated by tools/write_pyi.py - do not edit ### +# ### imports are manually managed + +from typing import Callable +from typing import ContextManager +from typing import Optional +from typing import TextIO +from typing import Tuple +from typing import TYPE_CHECKING +from typing import Union + +if TYPE_CHECKING: + from sqlalchemy.engine.base import Connection + from sqlalchemy.sql.schema import MetaData + + from .migration import MigrationContext + from .runtime.migration import _ProxyTransaction + +### end imports ### + +def begin_transaction() -> Union["_ProxyTransaction", ContextManager]: + """Return a context manager that will + enclose an operation within a "transaction", + as defined by the environment's offline + and transactional DDL settings. + + e.g.:: + + with context.begin_transaction(): + context.run_migrations() + + :meth:`.begin_transaction` is intended to + "do the right thing" regardless of + calling context: + + * If :meth:`.is_transactional_ddl` is ``False``, + returns a "do nothing" context manager + which otherwise produces no transactional + state or directives. + * If :meth:`.is_offline_mode` is ``True``, + returns a context manager that will + invoke the :meth:`.DefaultImpl.emit_begin` + and :meth:`.DefaultImpl.emit_commit` + methods, which will produce the string + directives ``BEGIN`` and ``COMMIT`` on + the output stream, as rendered by the + target backend (e.g. SQL Server would + emit ``BEGIN TRANSACTION``). + * Otherwise, calls :meth:`sqlalchemy.engine.Connection.begin` + on the current online connection, which + returns a :class:`sqlalchemy.engine.Transaction` + object. This object demarcates a real + transaction and is itself a context manager, + which will roll back if an exception + is raised. + + Note that a custom ``env.py`` script which + has more specific transactional needs can of course + manipulate the :class:`~sqlalchemy.engine.Connection` + directly to produce transactional state in "online" + mode. + + """ + +def configure( + connection: Optional["Connection"], + url: Optional[str], + dialect_name: Optional[str], + dialect_opts: Optional[dict], + transactional_ddl: Optional[bool], + transaction_per_migration: bool, + output_buffer: Optional[TextIO], + starting_rev: Optional[str], + tag: Optional[str], + template_args: Optional[dict], + render_as_batch: bool, + target_metadata: Optional["MetaData"], + include_name: Optional[Callable], + include_object: Optional[Callable], + include_schemas: bool, + process_revision_directives: Optional[Callable], + compare_type: bool, + compare_server_default: bool, + render_item: Optional[Callable], + literal_binds: bool, + upgrade_token: str, + downgrade_token: str, + alembic_module_prefix: str, + sqlalchemy_module_prefix: str, + user_module_prefix: Optional[str], + on_version_apply: Optional[Callable], + **kw +) -> None: + """Configure a :class:`.MigrationContext` within this + :class:`.EnvironmentContext` which will provide database + connectivity and other configuration to a series of + migration scripts. + + Many methods on :class:`.EnvironmentContext` require that + this method has been called in order to function, as they + ultimately need to have database access or at least access + to the dialect in use. Those which do are documented as such. + + The important thing needed by :meth:`.configure` is a + means to determine what kind of database dialect is in use. + An actual connection to that database is needed only if + the :class:`.MigrationContext` is to be used in + "online" mode. + + If the :meth:`.is_offline_mode` function returns ``True``, + then no connection is needed here. Otherwise, the + ``connection`` parameter should be present as an + instance of :class:`sqlalchemy.engine.Connection`. + + This function is typically called from the ``env.py`` + script within a migration environment. It can be called + multiple times for an invocation. The most recent + :class:`~sqlalchemy.engine.Connection` + for which it was called is the one that will be operated upon + by the next call to :meth:`.run_migrations`. + + General parameters: + + :param connection: a :class:`~sqlalchemy.engine.Connection` + to use + for SQL execution in "online" mode. When present, is also + used to determine the type of dialect in use. + :param url: a string database url, or a + :class:`sqlalchemy.engine.url.URL` object. + The type of dialect to be used will be derived from this if + ``connection`` is not passed. + :param dialect_name: string name of a dialect, such as + "postgresql", "mssql", etc. + The type of dialect to be used will be derived from this if + ``connection`` and ``url`` are not passed. + :param dialect_opts: dictionary of options to be passed to dialect + constructor. + + .. versionadded:: 1.0.12 + + :param transactional_ddl: Force the usage of "transactional" + DDL on or off; + this otherwise defaults to whether or not the dialect in + use supports it. + :param transaction_per_migration: if True, nest each migration script + in a transaction rather than the full series of migrations to + run. + :param output_buffer: a file-like object that will be used + for textual output + when the ``--sql`` option is used to generate SQL scripts. + Defaults to + ``sys.stdout`` if not passed here and also not present on + the :class:`.Config` + object. The value here overrides that of the :class:`.Config` + object. + :param output_encoding: when using ``--sql`` to generate SQL + scripts, apply this encoding to the string output. + :param literal_binds: when using ``--sql`` to generate SQL + scripts, pass through the ``literal_binds`` flag to the compiler + so that any literal values that would ordinarily be bound + parameters are converted to plain strings. + + .. warning:: Dialects can typically only handle simple datatypes + like strings and numbers for auto-literal generation. Datatypes + like dates, intervals, and others may still require manual + formatting, typically using :meth:`.Operations.inline_literal`. + + .. note:: the ``literal_binds`` flag is ignored on SQLAlchemy + versions prior to 0.8 where this feature is not supported. + + .. seealso:: + + :meth:`.Operations.inline_literal` + + :param starting_rev: Override the "starting revision" argument + when using ``--sql`` mode. + :param tag: a string tag for usage by custom ``env.py`` scripts. + Set via the ``--tag`` option, can be overridden here. + :param template_args: dictionary of template arguments which + will be added to the template argument environment when + running the "revision" command. Note that the script environment + is only run within the "revision" command if the --autogenerate + option is used, or if the option "revision_environment=true" + is present in the alembic.ini file. + + :param version_table: The name of the Alembic version table. + The default is ``'alembic_version'``. + :param version_table_schema: Optional schema to place version + table within. + :param version_table_pk: boolean, whether the Alembic version table + should use a primary key constraint for the "value" column; this + only takes effect when the table is first created. + Defaults to True; setting to False should not be necessary and is + here for backwards compatibility reasons. + :param on_version_apply: a callable or collection of callables to be + run for each migration step. + The callables will be run in the order they are given, once for + each migration step, after the respective operation has been + applied but before its transaction is finalized. + Each callable accepts no positional arguments and the following + keyword arguments: + + * ``ctx``: the :class:`.MigrationContext` running the migration, + * ``step``: a :class:`.MigrationInfo` representing the + step currently being applied, + * ``heads``: a collection of version strings representing the + current heads, + * ``run_args``: the ``**kwargs`` passed to :meth:`.run_migrations`. + + Parameters specific to the autogenerate feature, when + ``alembic revision`` is run with the ``--autogenerate`` feature: + + :param target_metadata: a :class:`sqlalchemy.schema.MetaData` + object, or a sequence of :class:`~sqlalchemy.schema.MetaData` + objects, that will be consulted during autogeneration. + The tables present in each :class:`~sqlalchemy.schema.MetaData` + will be compared against + what is locally available on the target + :class:`~sqlalchemy.engine.Connection` + to produce candidate upgrade/downgrade operations. + :param compare_type: Indicates type comparison behavior during + an autogenerate + operation. Defaults to ``False`` which disables type + comparison. Set to + ``True`` to turn on default type comparison, which has varied + accuracy depending on backend. See :ref:`compare_types` + for an example as well as information on other type + comparison options. + + .. seealso:: + + :ref:`compare_types` + + :paramref:`.EnvironmentContext.configure.compare_server_default` + + :param compare_server_default: Indicates server default comparison + behavior during + an autogenerate operation. Defaults to ``False`` which disables + server default + comparison. Set to ``True`` to turn on server default comparison, + which has + varied accuracy depending on backend. + + To customize server default comparison behavior, a callable may + be specified + which can filter server default comparisons during an + autogenerate operation. + defaults during an autogenerate operation. The format of this + callable is:: + + def my_compare_server_default(context, inspected_column, + metadata_column, inspected_default, metadata_default, + rendered_metadata_default): + # return True if the defaults are different, + # False if not, or None to allow the default implementation + # to compare these defaults + return None + + context.configure( + # ... + compare_server_default = my_compare_server_default + ) + + ``inspected_column`` is a dictionary structure as returned by + :meth:`sqlalchemy.engine.reflection.Inspector.get_columns`, whereas + ``metadata_column`` is a :class:`sqlalchemy.schema.Column` from + the local model environment. + + A return value of ``None`` indicates to allow default server default + comparison + to proceed. Note that some backends such as Postgresql actually + execute + the two defaults on the database side to compare for equivalence. + + .. seealso:: + + :paramref:`.EnvironmentContext.configure.compare_type` + + :param include_name: A callable function which is given + the chance to return ``True`` or ``False`` for any database reflected + object based on its name, including database schema names when + the :paramref:`.EnvironmentContext.configure.include_schemas` flag + is set to ``True``. + + The function accepts the following positional arguments: + + * ``name``: the name of the object, such as schema name or table name. + Will be ``None`` when indicating the default schema name of the + database connection. + * ``type``: a string describing the type of object; currently + ``"schema"``, ``"table"``, ``"column"``, ``"index"``, + ``"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"``. + + E.g.:: + + def include_name(name, type_, parent_names): + if type_ == "schema": + return name in ["schema_one", "schema_two"] + else: + return True + + context.configure( + # ... + include_schemas = True, + include_name = include_name + ) + + .. versionadded:: 1.5 + + .. seealso:: + + :ref:`autogenerate_include_hooks` + + :paramref:`.EnvironmentContext.configure.include_object` + + :paramref:`.EnvironmentContext.configure.include_schemas` + + + :param include_object: A callable function which is given + the chance to return ``True`` or ``False`` for any object, + indicating if the given object should be considered in the + autogenerate sweep. + + The function accepts the following positional arguments: + + * ``object``: a :class:`~sqlalchemy.schema.SchemaItem` object such + as a :class:`~sqlalchemy.schema.Table`, + :class:`~sqlalchemy.schema.Column`, + :class:`~sqlalchemy.schema.Index` + :class:`~sqlalchemy.schema.UniqueConstraint`, + or :class:`~sqlalchemy.schema.ForeignKeyConstraint` object + * ``name``: the name of the object. This is typically available + via ``object.name``. + * ``type``: a string describing the type of object; currently + ``"table"``, ``"column"``, ``"index"``, ``"unique_constraint"``, + or ``"foreign_key_constraint"`` + * ``reflected``: ``True`` if the given object was produced based on + table reflection, ``False`` if it's from a local :class:`.MetaData` + object. + * ``compare_to``: the object being compared against, if available, + else ``None``. + + E.g.:: + + def include_object(object, name, type_, reflected, compare_to): + if (type_ == "column" and + not reflected and + object.info.get("skip_autogenerate", False)): + return False + else: + return True + + context.configure( + # ... + include_object = include_object + ) + + For the use case of omitting specific schemas from a target database + when :paramref:`.EnvironmentContext.configure.include_schemas` is + set to ``True``, the :attr:`~sqlalchemy.schema.Table.schema` + attribute can be checked for each :class:`~sqlalchemy.schema.Table` + object passed to the hook, however it is much more efficient + to filter on schemas before reflection of objects takes place + using the :paramref:`.EnvironmentContext.configure.include_name` + hook. + + .. seealso:: + + :ref:`autogenerate_include_hooks` + + :paramref:`.EnvironmentContext.configure.include_name` + + :paramref:`.EnvironmentContext.configure.include_schemas` + + :param render_as_batch: if True, commands which alter elements + within a table will be placed under a ``with batch_alter_table():`` + directive, so that batch migrations will take place. + + .. seealso:: + + :ref:`batch_migrations` + + :param include_schemas: If True, autogenerate will scan across + all schemas located by the SQLAlchemy + :meth:`~sqlalchemy.engine.reflection.Inspector.get_schema_names` + method, and include all differences in tables found across all + those schemas. When using this option, you may want to also + use the :paramref:`.EnvironmentContext.configure.include_name` + parameter to specify a callable which + can filter the tables/schemas that get included. + + .. seealso:: + + :ref:`autogenerate_include_hooks` + + :paramref:`.EnvironmentContext.configure.include_name` + + :paramref:`.EnvironmentContext.configure.include_object` + + :param render_item: Callable that can be used to override how + any schema item, i.e. column, constraint, type, + etc., is rendered for autogenerate. The callable receives a + string describing the type of object, the object, and + the autogen context. If it returns False, the + default rendering method will be used. If it returns None, + the item will not be rendered in the context of a Table + construct, that is, can be used to skip columns or constraints + within op.create_table():: + + def my_render_column(type_, col, autogen_context): + if type_ == "column" and isinstance(col, MySpecialCol): + return repr(col) + else: + return False + + context.configure( + # ... + render_item = my_render_column + ) + + Available values for the type string include: ``"column"``, + ``"primary_key"``, ``"foreign_key"``, ``"unique"``, ``"check"``, + ``"type"``, ``"server_default"``. + + .. seealso:: + + :ref:`autogen_render_types` + + :param upgrade_token: When autogenerate completes, the text of the + candidate upgrade operations will be present in this template + variable when ``script.py.mako`` is rendered. Defaults to + ``upgrades``. + :param downgrade_token: When autogenerate completes, the text of the + candidate downgrade operations will be present in this + template variable when ``script.py.mako`` is rendered. Defaults to + ``downgrades``. + + :param alembic_module_prefix: When autogenerate refers to Alembic + :mod:`alembic.operations` constructs, this prefix will be used + (i.e. ``op.create_table``) Defaults to "``op.``". + Can be ``None`` to indicate no prefix. + + :param sqlalchemy_module_prefix: When autogenerate refers to + SQLAlchemy + :class:`~sqlalchemy.schema.Column` or type classes, this prefix + will be used + (i.e. ``sa.Column("somename", sa.Integer)``) Defaults to "``sa.``". + Can be ``None`` to indicate no prefix. + Note that when dialect-specific types are rendered, autogenerate + will render them using the dialect module name, i.e. ``mssql.BIT()``, + ``postgresql.UUID()``. + + :param user_module_prefix: When autogenerate refers to a SQLAlchemy + type (e.g. :class:`.TypeEngine`) where the module name is not + under the ``sqlalchemy`` namespace, this prefix will be used + within autogenerate. If left at its default of + ``None``, the ``__module__`` attribute of the type is used to + render the import module. It's a good practice to set this + and to have all custom types be available from a fixed module space, + in order to future-proof migration files against reorganizations + in modules. + + .. seealso:: + + :ref:`autogen_module_prefix` + + :param process_revision_directives: a callable function that will + be passed a structure representing the end result of an autogenerate + or plain "revision" operation, which can be manipulated to affect + how the ``alembic revision`` command ultimately outputs new + revision scripts. The structure of the callable is:: + + def process_revision_directives(context, revision, directives): + pass + + The ``directives`` parameter is a Python list containing + a single :class:`.MigrationScript` directive, which represents + the revision file to be generated. This list as well as its + contents may be freely modified to produce any set of commands. + The section :ref:`customizing_revision` shows an example of + doing this. The ``context`` parameter is the + :class:`.MigrationContext` in use, + and ``revision`` is a tuple of revision identifiers representing the + current revision of the database. + + The callable is invoked at all times when the ``--autogenerate`` + option is passed to ``alembic revision``. If ``--autogenerate`` + is not passed, the callable is invoked only if the + ``revision_environment`` variable is set to True in the Alembic + configuration, in which case the given ``directives`` collection + will contain empty :class:`.UpgradeOps` and :class:`.DowngradeOps` + collections for ``.upgrade_ops`` and ``.downgrade_ops``. The + ``--autogenerate`` option itself can be inferred by inspecting + ``context.config.cmd_opts.autogenerate``. + + The callable function may optionally be an instance of + a :class:`.Rewriter` object. This is a helper object that + assists in the production of autogenerate-stream rewriter functions. + + .. seealso:: + + :ref:`customizing_revision` + + :ref:`autogen_rewriter` + + :paramref:`.command.revision.process_revision_directives` + + Parameters specific to individual backends: + + :param mssql_batch_separator: The "batch separator" which will + be placed between each statement when generating offline SQL Server + migrations. Defaults to ``GO``. Note this is in addition to the + customary semicolon ``;`` at the end of each statement; SQL Server + considers the "batch separator" to denote the end of an + individual statement execution, and cannot group certain + dependent operations in one step. + :param oracle_batch_separator: The "batch separator" which will + be placed between each statement when generating offline + Oracle migrations. Defaults to ``/``. Oracle doesn't add a + semicolon between statements like most other backends. + + """ + +def execute(sql, execution_options): + """Execute the given SQL using the current change context. + + The behavior of :meth:`.execute` is the same + as that of :meth:`.Operations.execute`. Please see that + function's documentation for full detail including + caveats and limitations. + + This function requires that a :class:`.MigrationContext` has + first been made available via :meth:`.configure`. + + """ + +def get_bind(): + """Return the current 'bind'. + + In "online" mode, this is the + :class:`sqlalchemy.engine.Connection` currently being used + to emit SQL to the database. + + This function requires that a :class:`.MigrationContext` + has first been made available via :meth:`.configure`. + + """ + +def get_context() -> "MigrationContext": + """Return the current :class:`.MigrationContext` object. + + If :meth:`.EnvironmentContext.configure` has not been + called yet, raises an exception. + + """ + +def get_head_revision() -> Union[str, Tuple[str, ...], None]: + """Return the hex identifier of the 'head' script revision. + + If the script directory has multiple heads, this + method raises a :class:`.CommandError`; + :meth:`.EnvironmentContext.get_head_revisions` should be preferred. + + This function does not require that the :class:`.MigrationContext` + has been configured. + + .. seealso:: :meth:`.EnvironmentContext.get_head_revisions` + + """ + +def get_head_revisions() -> Union[str, Tuple[str, ...], None]: + """Return the hex identifier of the 'heads' script revision(s). + + This returns a tuple containing the version number of all + heads in the script directory. + + This function does not require that the :class:`.MigrationContext` + has been configured. + + """ + +def get_revision_argument() -> Union[str, Tuple[str, ...], None]: + """Get the 'destination' revision argument. + + This is typically the argument passed to the + ``upgrade`` or ``downgrade`` command. + + If it was specified as ``head``, the actual + version number is returned; if specified + as ``base``, ``None`` is returned. + + This function does not require that the :class:`.MigrationContext` + has been configured. + + """ + +def get_starting_revision_argument() -> Union[str, Tuple[str, ...], None]: + """Return the 'starting revision' argument, + if the revision was passed using ``start:end``. + + This is only meaningful in "offline" mode. + Returns ``None`` if no value is available + or was configured. + + This function does not require that the :class:`.MigrationContext` + has been configured. + + """ + +def get_tag_argument() -> Optional[str]: + """Return the value passed for the ``--tag`` argument, if any. + + The ``--tag`` argument is not used directly by Alembic, + but is available for custom ``env.py`` configurations that + wish to use it; particularly for offline generation scripts + that wish to generate tagged filenames. + + This function does not require that the :class:`.MigrationContext` + has been configured. + + .. seealso:: + + :meth:`.EnvironmentContext.get_x_argument` - a newer and more + open ended system of extending ``env.py`` scripts via the command + line. + + """ + +def get_x_argument(as_dictionary: bool): + """Return the value(s) passed for the ``-x`` argument, if any. + + The ``-x`` argument is an open ended flag that allows any user-defined + value or values to be passed on the command line, then available + here for consumption by a custom ``env.py`` script. + + The return value is a list, returned directly from the ``argparse`` + structure. If ``as_dictionary=True`` is passed, the ``x`` arguments + are parsed using ``key=value`` format into a dictionary that is + then returned. + + For example, to support passing a database URL on the command line, + the standard ``env.py`` script can be modified like this:: + + cmd_line_url = context.get_x_argument( + as_dictionary=True).get('dbname') + if cmd_line_url: + engine = create_engine(cmd_line_url) + else: + engine = engine_from_config( + config.get_section(config.config_ini_section), + prefix='sqlalchemy.', + poolclass=pool.NullPool) + + This then takes effect by running the ``alembic`` script as:: + + alembic -x dbname=postgresql://user:pass@host/dbname upgrade head + + This function does not require that the :class:`.MigrationContext` + has been configured. + + .. seealso:: + + :meth:`.EnvironmentContext.get_tag_argument` + + :attr:`.Config.cmd_opts` + + """ + +def is_offline_mode() -> bool: + """Return True if the current migrations environment + is running in "offline mode". + + This is ``True`` or ``False`` depending + on the ``--sql`` flag passed. + + This function does not require that the :class:`.MigrationContext` + has been configured. + + """ + +def is_transactional_ddl(): + """Return True if the context is configured to expect a + transactional DDL capable backend. + + This defaults to the type of database in use, and + can be overridden by the ``transactional_ddl`` argument + to :meth:`.configure` + + This function requires that a :class:`.MigrationContext` + has first been made available via :meth:`.configure`. + + """ + +def run_migrations(**kw) -> None: + """Run migrations as determined by the current command line + configuration + as well as versioning information present (or not) in the current + database connection (if one is present). + + The function accepts optional ``**kw`` arguments. If these are + passed, they are sent directly to the ``upgrade()`` and + ``downgrade()`` + functions within each target revision file. By modifying the + ``script.py.mako`` file so that the ``upgrade()`` and ``downgrade()`` + functions accept arguments, parameters can be passed here so that + contextual information, usually information to identify a particular + database in use, can be passed from a custom ``env.py`` script + to the migration functions. + + This function requires that a :class:`.MigrationContext` has + first been made available via :meth:`.configure`. + + """ + +def static_output(text): + """Emit text directly to the "offline" SQL stream. + + Typically this is for emitting comments that + start with --. The statement is not treated + as a SQL execution, no ; or batch separator + is added, etc. + + """ diff --git a/alembic/op.pyi b/alembic/op.pyi new file mode 100644 index 0000000..76912dc --- /dev/null +++ b/alembic/op.pyi @@ -0,0 +1,1160 @@ +# ### this file stubs are generated by tools/write_pyi.py - do not edit ### +# ### imports are manually managed + +from typing import Any +from typing import Callable +from typing import List +from typing import Optional +from typing import Sequence +from typing import Type +from typing import TYPE_CHECKING +from typing import Union + +from sqlalchemy.sql.expression import TableClause +from sqlalchemy.sql.expression import Update + +if TYPE_CHECKING: + + from sqlalchemy.engine import Connection + from sqlalchemy.sql.elements import BinaryExpression + from sqlalchemy.sql.elements import conv + from sqlalchemy.sql.elements import TextClause + from sqlalchemy.sql.functions import Function + from sqlalchemy.sql.schema import Column + from sqlalchemy.sql.schema import Computed + from sqlalchemy.sql.schema import Identity + from sqlalchemy.sql.schema import Table + from sqlalchemy.sql.type_api import TypeEngine + + from .operations.ops import MigrateOperation + from .util.sqla_compat import _literal_bindparam + +### end imports ### + +def add_column( + table_name: str, column: "Column", schema: Optional[str] +) -> Optional["Table"]: + """Issue an "add column" instruction using the current + migration context. + + e.g.:: + + from alembic import op + from sqlalchemy import Column, String + + op.add_column('organization', + Column('name', String()) + ) + + The provided :class:`~sqlalchemy.schema.Column` object can also + specify a :class:`~sqlalchemy.schema.ForeignKey`, referencing + a remote table name. Alembic will automatically generate a stub + "referenced" table and emit a second ALTER statement in order + to add the constraint separately:: + + from alembic import op + from sqlalchemy import Column, INTEGER, ForeignKey + + op.add_column('organization', + Column('account_id', INTEGER, ForeignKey('accounts.id')) + ) + + Note that this statement uses the :class:`~sqlalchemy.schema.Column` + construct as is from the SQLAlchemy library. In particular, + default values to be created on the database side are + specified using the ``server_default`` parameter, and not + ``default`` which only specifies Python-side defaults:: + + from alembic import op + from sqlalchemy import Column, TIMESTAMP, func + + # specify "DEFAULT NOW" along with the column add + op.add_column('account', + Column('timestamp', TIMESTAMP, server_default=func.now()) + ) + + :param table_name: String name of the parent table. + :param column: a :class:`sqlalchemy.schema.Column` object + representing the new column. + :param schema: Optional schema name to operate within. To control + quoting of the schema outside of the default behavior, use + the SQLAlchemy construct + :class:`~sqlalchemy.sql.elements.quoted_name`. + + """ + +def alter_column( + table_name: str, + column_name: str, + nullable: Optional[bool], + comment: Union[str, bool, None], + server_default: Any, + new_column_name: Optional[str], + type_: Union["TypeEngine", Type["TypeEngine"], None], + existing_type: Union["TypeEngine", Type["TypeEngine"], None], + existing_server_default: Union[str, bool, "Identity", "Computed", None], + existing_nullable: Optional[bool], + existing_comment: Optional[str], + schema: Optional[str], + **kw +) -> Optional["Table"]: + """Issue an "alter column" instruction using the + current migration context. + + Generally, only that aspect of the column which + is being changed, i.e. name, type, nullability, + default, needs to be specified. Multiple changes + can also be specified at once and the backend should + "do the right thing", emitting each change either + separately or together as the backend allows. + + MySQL has special requirements here, since MySQL + cannot ALTER a column without a full specification. + When producing MySQL-compatible migration files, + it is recommended that the ``existing_type``, + ``existing_server_default``, and ``existing_nullable`` + parameters be present, if not being altered. + + Type changes which are against the SQLAlchemy + "schema" types :class:`~sqlalchemy.types.Boolean` + and :class:`~sqlalchemy.types.Enum` may also + add or drop constraints which accompany those + types on backends that don't support them natively. + The ``existing_type`` argument is + used in this case to identify and remove a previous + constraint that was bound to the type object. + + :param table_name: string name of the target table. + :param column_name: string name of the target column, + as it exists before the operation begins. + :param nullable: Optional; specify ``True`` or ``False`` + to alter the column's nullability. + :param server_default: Optional; specify a string + SQL expression, :func:`~sqlalchemy.sql.expression.text`, + or :class:`~sqlalchemy.schema.DefaultClause` to indicate + an alteration to the column's default value. + Set to ``None`` to have the default removed. + :param comment: optional string text of a new comment to add to the + column. + + .. versionadded:: 1.0.6 + + :param new_column_name: Optional; specify a string name here to + indicate the new name within a column rename operation. + :param type\_: Optional; a :class:`~sqlalchemy.types.TypeEngine` + type object to specify a change to the column's type. + For SQLAlchemy types that also indicate a constraint (i.e. + :class:`~sqlalchemy.types.Boolean`, :class:`~sqlalchemy.types.Enum`), + the constraint is also generated. + :param autoincrement: set the ``AUTO_INCREMENT`` flag of the column; + currently understood by the MySQL dialect. + :param existing_type: Optional; a + :class:`~sqlalchemy.types.TypeEngine` + type object to specify the previous type. This + is required for all MySQL column alter operations that + don't otherwise specify a new type, as well as for + when nullability is being changed on a SQL Server + column. It is also used if the type is a so-called + SQLlchemy "schema" type which may define a constraint (i.e. + :class:`~sqlalchemy.types.Boolean`, + :class:`~sqlalchemy.types.Enum`), + so that the constraint can be dropped. + :param existing_server_default: Optional; The existing + default value of the column. Required on MySQL if + an existing default is not being changed; else MySQL + removes the default. + :param existing_nullable: Optional; the existing nullability + of the column. Required on MySQL if the existing nullability + is not being changed; else MySQL sets this to NULL. + :param existing_autoincrement: Optional; the existing autoincrement + of the column. Used for MySQL's system of altering a column + that specifies ``AUTO_INCREMENT``. + :param existing_comment: string text of the existing comment on the + column to be maintained. Required on MySQL if the existing comment + on the column is not being changed. + + .. versionadded:: 1.0.6 + + :param schema: Optional schema name to operate within. To control + quoting of the schema outside of the default behavior, use + the SQLAlchemy construct + :class:`~sqlalchemy.sql.elements.quoted_name`. + :param postgresql_using: String argument which will indicate a + SQL expression to render within the Postgresql-specific USING clause + within ALTER COLUMN. This string is taken directly as raw SQL which + must explicitly include any necessary quoting or escaping of tokens + within the expression. + + """ + +def batch_alter_table( + table_name, + schema, + recreate, + partial_reordering, + copy_from, + table_args, + table_kwargs, + reflect_args, + reflect_kwargs, + naming_convention, +): + """Invoke a series of per-table migrations in batch. + + Batch mode allows a series of operations specific to a table + to be syntactically grouped together, and allows for alternate + modes of table migration, in particular the "recreate" style of + migration required by SQLite. + + "recreate" style is as follows: + + 1. A new table is created with the new specification, based on the + migration directives within the batch, using a temporary name. + + 2. the data copied from the existing table to the new table. + + 3. the existing table is dropped. + + 4. the new table is renamed to the existing table name. + + The directive by default will only use "recreate" style on the + SQLite backend, and only if directives are present which require + this form, e.g. anything other than ``add_column()``. The batch + operation on other backends will proceed using standard ALTER TABLE + operations. + + The method is used as a context manager, which returns an instance + of :class:`.BatchOperations`; this object is the same as + :class:`.Operations` except that table names and schema names + are omitted. E.g.:: + + with op.batch_alter_table("some_table") as batch_op: + batch_op.add_column(Column('foo', Integer)) + batch_op.drop_column('bar') + + The operations within the context manager are invoked at once + when the context is ended. When run against SQLite, if the + migrations include operations not supported by SQLite's ALTER TABLE, + the entire table will be copied to a new one with the new + specification, moving all data across as well. + + The copy operation by default uses reflection to retrieve the current + structure of the table, and therefore :meth:`.batch_alter_table` + in this mode requires that the migration is run in "online" mode. + The ``copy_from`` parameter may be passed which refers to an existing + :class:`.Table` object, which will bypass this reflection step. + + .. note:: The table copy operation will currently not copy + CHECK constraints, and may not copy UNIQUE constraints that are + unnamed, as is possible on SQLite. See the section + :ref:`sqlite_batch_constraints` for workarounds. + + :param table_name: name of table + :param schema: optional schema name. + :param recreate: under what circumstances the table should be + recreated. At its default of ``"auto"``, the SQLite dialect will + recreate the table if any operations other than ``add_column()``, + ``create_index()``, or ``drop_index()`` are + present. Other options include ``"always"`` and ``"never"``. + :param copy_from: optional :class:`~sqlalchemy.schema.Table` object + that will act as the structure of the table being copied. If omitted, + table reflection is used to retrieve the structure of the table. + + .. seealso:: + + :ref:`batch_offline_mode` + + :paramref:`~.Operations.batch_alter_table.reflect_args` + + :paramref:`~.Operations.batch_alter_table.reflect_kwargs` + + :param reflect_args: a sequence of additional positional arguments that + will be applied to the table structure being reflected / copied; + this may be used to pass column and constraint overrides to the + table that will be reflected, in lieu of passing the whole + :class:`~sqlalchemy.schema.Table` using + :paramref:`~.Operations.batch_alter_table.copy_from`. + :param reflect_kwargs: a dictionary of additional keyword arguments + that will be applied to the table structure being copied; this may be + used to pass additional table and reflection options to the table that + will be reflected, in lieu of passing the whole + :class:`~sqlalchemy.schema.Table` using + :paramref:`~.Operations.batch_alter_table.copy_from`. + :param table_args: a sequence of additional positional arguments that + will be applied to the new :class:`~sqlalchemy.schema.Table` when + created, in addition to those copied from the source table. + This may be used to provide additional constraints such as CHECK + constraints that may not be reflected. + :param table_kwargs: a dictionary of additional keyword arguments + that will be applied to the new :class:`~sqlalchemy.schema.Table` + when created, in addition to those copied from the source table. + This may be used to provide for additional table options that may + not be reflected. + :param naming_convention: a naming convention dictionary of the form + described at :ref:`autogen_naming_conventions` which will be applied + to the :class:`~sqlalchemy.schema.MetaData` during the reflection + process. This is typically required if one wants to drop SQLite + constraints, as these constraints will not have names when + reflected on this backend. Requires SQLAlchemy **0.9.4** or greater. + + .. seealso:: + + :ref:`dropping_sqlite_foreign_keys` + + :param partial_reordering: a list of tuples, each suggesting a desired + ordering of two or more columns in the newly created table. Requires + that :paramref:`.batch_alter_table.recreate` is set to ``"always"``. + Examples, given a table with columns "a", "b", "c", and "d": + + Specify the order of all columns:: + + with op.batch_alter_table( + "some_table", recreate="always", + partial_reordering=[("c", "d", "a", "b")] + ) as batch_op: + pass + + Ensure "d" appears before "c", and "b", appears before "a":: + + with op.batch_alter_table( + "some_table", recreate="always", + partial_reordering=[("d", "c"), ("b", "a")] + ) as batch_op: + pass + + The ordering of columns not included in the partial_reordering + set is undefined. Therefore it is best to specify the complete + ordering of all columns for best results. + + .. versionadded:: 1.4.0 + + .. note:: batch mode requires SQLAlchemy 0.8 or above. + + .. seealso:: + + :ref:`batch_migrations` + + """ + +def bulk_insert( + table: Union["Table", "TableClause"], rows: List[dict], multiinsert: bool +) -> None: + """Issue a "bulk insert" operation using the current + migration context. + + This provides a means of representing an INSERT of multiple rows + which works equally well in the context of executing on a live + connection as well as that of generating a SQL script. In the + case of a SQL script, the values are rendered inline into the + statement. + + e.g.:: + + from alembic import op + from datetime import date + from sqlalchemy.sql import table, column + from sqlalchemy import String, Integer, Date + + # Create an ad-hoc table to use for the insert statement. + accounts_table = table('account', + column('id', Integer), + column('name', String), + column('create_date', Date) + ) + + op.bulk_insert(accounts_table, + [ + {'id':1, 'name':'John Smith', + 'create_date':date(2010, 10, 5)}, + {'id':2, 'name':'Ed Williams', + 'create_date':date(2007, 5, 27)}, + {'id':3, 'name':'Wendy Jones', + 'create_date':date(2008, 8, 15)}, + ] + ) + + When using --sql mode, some datatypes may not render inline + automatically, such as dates and other special types. When this + issue is present, :meth:`.Operations.inline_literal` may be used:: + + op.bulk_insert(accounts_table, + [ + {'id':1, 'name':'John Smith', + 'create_date':op.inline_literal("2010-10-05")}, + {'id':2, 'name':'Ed Williams', + 'create_date':op.inline_literal("2007-05-27")}, + {'id':3, 'name':'Wendy Jones', + 'create_date':op.inline_literal("2008-08-15")}, + ], + multiinsert=False + ) + + When using :meth:`.Operations.inline_literal` in conjunction with + :meth:`.Operations.bulk_insert`, in order for the statement to work + in "online" (e.g. non --sql) mode, the + :paramref:`~.Operations.bulk_insert.multiinsert` + flag should be set to ``False``, which will have the effect of + individual INSERT statements being emitted to the database, each + with a distinct VALUES clause, so that the "inline" values can + still be rendered, rather than attempting to pass the values + as bound parameters. + + :param table: a table object which represents the target of the INSERT. + + :param rows: a list of dictionaries indicating rows. + + :param multiinsert: when at its default of True and --sql mode is not + enabled, the INSERT statement will be executed using + "executemany()" style, where all elements in the list of + dictionaries are passed as bound parameters in a single + list. Setting this to False results in individual INSERT + statements being emitted per parameter set, and is needed + in those cases where non-literal values are present in the + parameter sets. + + """ + +def create_check_constraint( + constraint_name: Optional[str], + table_name: str, + condition: "BinaryExpression", + schema: Optional[str], + **kw +) -> Optional["Table"]: + """Issue a "create check constraint" instruction using the + current migration context. + + e.g.:: + + from alembic import op + from sqlalchemy.sql import column, func + + op.create_check_constraint( + "ck_user_name_len", + "user", + func.len(column('name')) > 5 + ) + + CHECK constraints are usually against a SQL expression, so ad-hoc + table metadata is usually needed. The function will convert the given + arguments into a :class:`sqlalchemy.schema.CheckConstraint` bound + to an anonymous table in order to emit the CREATE statement. + + :param name: Name of the check constraint. The name is necessary + so that an ALTER statement can be emitted. For setups that + use an automated naming scheme such as that described at + :ref:`sqla:constraint_naming_conventions`, + ``name`` here can be ``None``, as the event listener will + apply the name to the constraint object when it is associated + with the table. + :param table_name: String name of the source table. + :param condition: SQL expression that's the condition of the + constraint. Can be a string or SQLAlchemy expression language + structure. + :param deferrable: optional bool. If set, emit DEFERRABLE or + NOT DEFERRABLE when issuing DDL for this constraint. + :param initially: optional string. If set, emit INITIALLY <value> + when issuing DDL for this constraint. + :param schema: Optional schema name to operate within. To control + quoting of the schema outside of the default behavior, use + the SQLAlchemy construct + :class:`~sqlalchemy.sql.elements.quoted_name`. + + """ + +def create_exclude_constraint( + constraint_name: str, table_name: str, *elements: Any, **kw: Any +) -> Optional["Table"]: + """Issue an alter to create an EXCLUDE constraint using the + current migration context. + + .. note:: This method is Postgresql specific, and additionally + requires at least SQLAlchemy 1.0. + + e.g.:: + + from alembic import op + + op.create_exclude_constraint( + "user_excl", + "user", + + ("period", '&&'), + ("group", '='), + where=("group != 'some group'") + + ) + + Note that the expressions work the same way as that of + the ``ExcludeConstraint`` object itself; if plain strings are + passed, quoting rules must be applied manually. + + :param name: Name of the constraint. + :param table_name: String name of the source table. + :param elements: exclude conditions. + :param where: SQL expression or SQL string with optional WHERE + clause. + :param deferrable: optional bool. If set, emit DEFERRABLE or + NOT DEFERRABLE when issuing DDL for this constraint. + :param initially: optional string. If set, emit INITIALLY <value> + when issuing DDL for this constraint. + :param schema: Optional schema name to operate within. + + """ + +def create_foreign_key( + constraint_name: str, + source_table: str, + referent_table: str, + local_cols: List[str], + remote_cols: List[str], + onupdate: Optional[str], + ondelete: Optional[str], + deferrable: Optional[bool], + initially: Optional[str], + match: Optional[str], + source_schema: Optional[str], + referent_schema: Optional[str], + **dialect_kw +) -> Optional["Table"]: + """Issue a "create foreign key" instruction using the + current migration context. + + e.g.:: + + from alembic import op + op.create_foreign_key( + "fk_user_address", "address", + "user", ["user_id"], ["id"]) + + This internally generates a :class:`~sqlalchemy.schema.Table` object + containing the necessary columns, then generates a new + :class:`~sqlalchemy.schema.ForeignKeyConstraint` + object which it then associates with the + :class:`~sqlalchemy.schema.Table`. + Any event listeners associated with this action will be fired + off normally. The :class:`~sqlalchemy.schema.AddConstraint` + construct is ultimately used to generate the ALTER statement. + + :param name: Name of the foreign key constraint. The name is necessary + so that an ALTER statement can be emitted. For setups that + use an automated naming scheme such as that described at + :ref:`sqla:constraint_naming_conventions`, + ``name`` here can be ``None``, as the event listener will + apply the name to the constraint object when it is associated + with the table. + :param source_table: String name of the source table. + :param referent_table: String name of the destination table. + :param local_cols: a list of string column names in the + source table. + :param remote_cols: a list of string column names in the + remote table. + :param onupdate: Optional string. If set, emit ON UPDATE <value> when + issuing DDL for this constraint. Typical values include CASCADE, + DELETE and RESTRICT. + :param ondelete: Optional string. If set, emit ON DELETE <value> when + issuing DDL for this constraint. Typical values include CASCADE, + DELETE and RESTRICT. + :param deferrable: optional bool. If set, emit DEFERRABLE or NOT + DEFERRABLE when issuing DDL for this constraint. + :param source_schema: Optional schema name of the source table. + :param referent_schema: Optional schema name of the destination table. + + """ + +def create_index( + index_name: str, + table_name: str, + columns: Sequence[Union[str, "TextClause", "Function"]], + schema: Optional[str], + unique: bool, + **kw +) -> Optional["Table"]: + """Issue a "create index" instruction using the current + migration context. + + e.g.:: + + from alembic import op + op.create_index('ik_test', 't1', ['foo', 'bar']) + + Functional indexes can be produced by using the + :func:`sqlalchemy.sql.expression.text` construct:: + + from alembic import op + from sqlalchemy import text + op.create_index('ik_test', 't1', [text('lower(foo)')]) + + :param index_name: name of the index. + :param table_name: name of the owning table. + :param columns: a list consisting of string column names and/or + :func:`~sqlalchemy.sql.expression.text` constructs. + :param schema: Optional schema name to operate within. To control + quoting of the schema outside of the default behavior, use + the SQLAlchemy construct + :class:`~sqlalchemy.sql.elements.quoted_name`. + :param unique: If True, create a unique index. + + :param quote: + Force quoting of this column's name on or off, corresponding + to ``True`` or ``False``. When left at its default + of ``None``, the column identifier will be quoted according to + whether the name is case sensitive (identifiers with at least one + upper case character are treated as case sensitive), or if it's a + reserved word. This flag is only needed to force quoting of a + reserved word which is not known by the SQLAlchemy dialect. + + :param \**kw: Additional keyword arguments not mentioned above are + dialect specific, and passed in the form + ``<dialectname>_<argname>``. + See the documentation regarding an individual dialect at + :ref:`dialect_toplevel` for detail on documented arguments. + + """ + +def create_primary_key( + constraint_name: str, + table_name: str, + columns: List[str], + schema: Optional[str], +) -> Optional["Table"]: + """Issue a "create primary key" instruction using the current + migration context. + + e.g.:: + + from alembic import op + op.create_primary_key( + "pk_my_table", "my_table", + ["id", "version"] + ) + + This internally generates a :class:`~sqlalchemy.schema.Table` object + containing the necessary columns, then generates a new + :class:`~sqlalchemy.schema.PrimaryKeyConstraint` + object which it then associates with the + :class:`~sqlalchemy.schema.Table`. + Any event listeners associated with this action will be fired + off normally. The :class:`~sqlalchemy.schema.AddConstraint` + construct is ultimately used to generate the ALTER statement. + + :param name: Name of the primary key constraint. The name is necessary + so that an ALTER statement can be emitted. For setups that + use an automated naming scheme such as that described at + :ref:`sqla:constraint_naming_conventions` + ``name`` here can be ``None``, as the event listener will + apply the name to the constraint object when it is associated + with the table. + :param table_name: String name of the target table. + :param columns: a list of string column names to be applied to the + primary key constraint. + :param schema: Optional schema name to operate within. To control + quoting of the schema outside of the default behavior, use + the SQLAlchemy construct + :class:`~sqlalchemy.sql.elements.quoted_name`. + + """ + +def create_table(table_name: str, *columns, **kw) -> Optional["Table"]: + """Issue a "create table" instruction using the current migration + context. + + This directive receives an argument list similar to that of the + traditional :class:`sqlalchemy.schema.Table` construct, but without the + metadata:: + + from sqlalchemy import INTEGER, VARCHAR, NVARCHAR, Column + from alembic import op + + op.create_table( + 'account', + Column('id', INTEGER, primary_key=True), + Column('name', VARCHAR(50), nullable=False), + Column('description', NVARCHAR(200)), + Column('timestamp', TIMESTAMP, server_default=func.now()) + ) + + Note that :meth:`.create_table` accepts + :class:`~sqlalchemy.schema.Column` + constructs directly from the SQLAlchemy library. In particular, + default values to be created on the database side are + specified using the ``server_default`` parameter, and not + ``default`` which only specifies Python-side defaults:: + + from alembic import op + from sqlalchemy import Column, TIMESTAMP, func + + # specify "DEFAULT NOW" along with the "timestamp" column + op.create_table('account', + Column('id', INTEGER, primary_key=True), + Column('timestamp', TIMESTAMP, server_default=func.now()) + ) + + The function also returns a newly created + :class:`~sqlalchemy.schema.Table` object, corresponding to the table + specification given, which is suitable for + immediate SQL operations, in particular + :meth:`.Operations.bulk_insert`:: + + from sqlalchemy import INTEGER, VARCHAR, NVARCHAR, Column + from alembic import op + + account_table = op.create_table( + 'account', + Column('id', INTEGER, primary_key=True), + Column('name', VARCHAR(50), nullable=False), + Column('description', NVARCHAR(200)), + Column('timestamp', TIMESTAMP, server_default=func.now()) + ) + + op.bulk_insert( + account_table, + [ + {"name": "A1", "description": "account 1"}, + {"name": "A2", "description": "account 2"}, + ] + ) + + :param table_name: Name of the table + :param \*columns: collection of :class:`~sqlalchemy.schema.Column` + objects within + the table, as well as optional :class:`~sqlalchemy.schema.Constraint` + objects + and :class:`~.sqlalchemy.schema.Index` objects. + :param schema: Optional schema name to operate within. To control + quoting of the schema outside of the default behavior, use + the SQLAlchemy construct + :class:`~sqlalchemy.sql.elements.quoted_name`. + :param \**kw: Other keyword arguments are passed to the underlying + :class:`sqlalchemy.schema.Table` object created for the command. + + :return: the :class:`~sqlalchemy.schema.Table` object corresponding + to the parameters given. + + """ + +def create_table_comment( + table_name: str, + comment: Optional[str], + existing_comment: None, + schema: Optional[str], +) -> Optional["Table"]: + """Emit a COMMENT ON operation to set the comment for a table. + + .. versionadded:: 1.0.6 + + :param table_name: string name of the target table. + :param comment: string value of the comment being registered against + the specified table. + :param existing_comment: String value of a comment + already registered on the specified table, used within autogenerate + so that the operation is reversible, but not required for direct + use. + + .. seealso:: + + :meth:`.Operations.drop_table_comment` + + :paramref:`.Operations.alter_column.comment` + + """ + +def create_unique_constraint( + constraint_name: Optional[str], + table_name: str, + columns: Sequence[str], + schema: Optional[str], + **kw +) -> Any: + """Issue a "create unique constraint" instruction using the + current migration context. + + e.g.:: + + from alembic import op + op.create_unique_constraint("uq_user_name", "user", ["name"]) + + This internally generates a :class:`~sqlalchemy.schema.Table` object + containing the necessary columns, then generates a new + :class:`~sqlalchemy.schema.UniqueConstraint` + object which it then associates with the + :class:`~sqlalchemy.schema.Table`. + Any event listeners associated with this action will be fired + off normally. The :class:`~sqlalchemy.schema.AddConstraint` + construct is ultimately used to generate the ALTER statement. + + :param name: Name of the unique constraint. The name is necessary + so that an ALTER statement can be emitted. For setups that + use an automated naming scheme such as that described at + :ref:`sqla:constraint_naming_conventions`, + ``name`` here can be ``None``, as the event listener will + apply the name to the constraint object when it is associated + with the table. + :param table_name: String name of the source table. + :param columns: a list of string column names in the + source table. + :param deferrable: optional bool. If set, emit DEFERRABLE or + NOT DEFERRABLE when issuing DDL for this constraint. + :param initially: optional string. If set, emit INITIALLY <value> + when issuing DDL for this constraint. + :param schema: Optional schema name to operate within. To control + quoting of the schema outside of the default behavior, use + the SQLAlchemy construct + :class:`~sqlalchemy.sql.elements.quoted_name`. + + """ + +def drop_column( + table_name: str, column_name: str, schema: Optional[str], **kw +) -> Optional["Table"]: + """Issue a "drop column" instruction using the current + migration context. + + e.g.:: + + drop_column('organization', 'account_id') + + :param table_name: name of table + :param column_name: name of column + :param schema: Optional schema name to operate within. To control + quoting of the schema outside of the default behavior, use + the SQLAlchemy construct + :class:`~sqlalchemy.sql.elements.quoted_name`. + :param mssql_drop_check: Optional boolean. When ``True``, on + Microsoft SQL Server only, first + drop the CHECK constraint on the column using a + SQL-script-compatible + block that selects into a @variable from sys.check_constraints, + then exec's a separate DROP CONSTRAINT for that constraint. + :param mssql_drop_default: Optional boolean. When ``True``, on + Microsoft SQL Server only, first + drop the DEFAULT constraint on the column using a + SQL-script-compatible + block that selects into a @variable from sys.default_constraints, + then exec's a separate DROP CONSTRAINT for that default. + :param mssql_drop_foreign_key: Optional boolean. When ``True``, on + Microsoft SQL Server only, first + drop a single FOREIGN KEY constraint on the column using a + SQL-script-compatible + block that selects into a @variable from + sys.foreign_keys/sys.foreign_key_columns, + then exec's a separate DROP CONSTRAINT for that default. Only + works if the column has exactly one FK constraint which refers to + it, at the moment. + + """ + +def drop_constraint( + constraint_name: str, + table_name: str, + type_: Optional[str], + schema: Optional[str], +) -> Optional["Table"]: + """Drop a constraint of the given name, typically via DROP CONSTRAINT. + + :param constraint_name: name of the constraint. + :param table_name: table name. + :param type\_: optional, required on MySQL. can be + 'foreignkey', 'primary', 'unique', or 'check'. + :param schema: Optional schema name to operate within. To control + quoting of the schema outside of the default behavior, use + the SQLAlchemy construct + :class:`~sqlalchemy.sql.elements.quoted_name`. + + """ + +def drop_index( + index_name: str, table_name: Optional[str], schema: Optional[str], **kw +) -> Optional["Table"]: + """Issue a "drop index" instruction using the current + migration context. + + e.g.:: + + drop_index("accounts") + + :param index_name: name of the index. + :param table_name: name of the owning table. Some + backends such as Microsoft SQL Server require this. + :param schema: Optional schema name to operate within. To control + quoting of the schema outside of the default behavior, use + the SQLAlchemy construct + :class:`~sqlalchemy.sql.elements.quoted_name`. + :param \**kw: Additional keyword arguments not mentioned above are + dialect specific, and passed in the form + ``<dialectname>_<argname>``. + See the documentation regarding an individual dialect at + :ref:`dialect_toplevel` for detail on documented arguments. + + """ + +def drop_table(table_name: str, schema: Optional[str], **kw: Any) -> None: + """Issue a "drop table" instruction using the current + migration context. + + + e.g.:: + + drop_table("accounts") + + :param table_name: Name of the table + :param schema: Optional schema name to operate within. To control + quoting of the schema outside of the default behavior, use + the SQLAlchemy construct + :class:`~sqlalchemy.sql.elements.quoted_name`. + :param \**kw: Other keyword arguments are passed to the underlying + :class:`sqlalchemy.schema.Table` object created for the command. + + """ + +def drop_table_comment( + table_name: str, existing_comment: Optional[str], schema: Optional[str] +) -> Optional["Table"]: + """Issue a "drop table comment" operation to + remove an existing comment set on a table. + + .. versionadded:: 1.0.6 + + :param table_name: string name of the target table. + :param existing_comment: An optional string value of a comment already + registered on the specified table. + + .. seealso:: + + :meth:`.Operations.create_table_comment` + + :paramref:`.Operations.alter_column.comment` + + """ + +def execute( + sqltext: Union[str, "TextClause", "Update"], execution_options: None +) -> Optional["Table"]: + """Execute the given SQL using the current migration context. + + The given SQL can be a plain string, e.g.:: + + op.execute("INSERT INTO table (foo) VALUES ('some value')") + + Or it can be any kind of Core SQL Expression construct, such as + below where we use an update construct:: + + from sqlalchemy.sql import table, column + from sqlalchemy import String + from alembic import op + + account = table('account', + column('name', String) + ) + op.execute( + account.update().\\ + where(account.c.name==op.inline_literal('account 1')).\\ + values({'name':op.inline_literal('account 2')}) + ) + + Above, we made use of the SQLAlchemy + :func:`sqlalchemy.sql.expression.table` and + :func:`sqlalchemy.sql.expression.column` constructs to make a brief, + ad-hoc table construct just for our UPDATE statement. A full + :class:`~sqlalchemy.schema.Table` construct of course works perfectly + fine as well, though note it's a recommended practice to at least + ensure the definition of a table is self-contained within the migration + script, rather than imported from a module that may break compatibility + with older migrations. + + In a SQL script context, the statement is emitted directly to the + output stream. There is *no* return result, however, as this + function is oriented towards generating a change script + that can run in "offline" mode. Additionally, parameterized + statements are discouraged here, as they *will not work* in offline + mode. Above, we use :meth:`.inline_literal` where parameters are + to be used. + + For full interaction with a connected database where parameters can + also be used normally, use the "bind" available from the context:: + + from alembic import op + connection = op.get_bind() + + connection.execute( + account.update().where(account.c.name=='account 1'). + values({"name": "account 2"}) + ) + + Additionally, when passing the statement as a plain string, it is first + coerceed into a :func:`sqlalchemy.sql.expression.text` construct + before being passed along. In the less likely case that the + literal SQL string contains a colon, it must be escaped with a + backslash, as:: + + op.execute("INSERT INTO table (foo) VALUES ('\:colon_value')") + + + :param sql: Any legal SQLAlchemy expression, including: + + * a string + * a :func:`sqlalchemy.sql.expression.text` construct. + * a :func:`sqlalchemy.sql.expression.insert` construct. + * a :func:`sqlalchemy.sql.expression.update`, + :func:`sqlalchemy.sql.expression.insert`, + or :func:`sqlalchemy.sql.expression.delete` construct. + * Pretty much anything that's "executable" as described + in :ref:`sqlexpression_toplevel`. + + .. note:: when passing a plain string, the statement is coerced into + a :func:`sqlalchemy.sql.expression.text` construct. This construct + considers symbols with colons, e.g. ``:foo`` to be bound parameters. + To avoid this, ensure that colon symbols are escaped, e.g. + ``\:foo``. + + :param execution_options: Optional dictionary of + execution options, will be passed to + :meth:`sqlalchemy.engine.Connection.execution_options`. + """ + +def f(name: str) -> "conv": + """Indicate a string name that has already had a naming convention + applied to it. + + This feature combines with the SQLAlchemy ``naming_convention`` feature + to disambiguate constraint names that have already had naming + conventions applied to them, versus those that have not. This is + necessary in the case that the ``"%(constraint_name)s"`` token + is used within a naming convention, so that it can be identified + that this particular name should remain fixed. + + If the :meth:`.Operations.f` is used on a constraint, the naming + convention will not take effect:: + + op.add_column('t', 'x', Boolean(name=op.f('ck_bool_t_x'))) + + Above, the CHECK constraint generated will have the name + ``ck_bool_t_x`` regardless of whether or not a naming convention is + in use. + + Alternatively, if a naming convention is in use, and 'f' is not used, + names will be converted along conventions. If the ``target_metadata`` + contains the naming convention + ``{"ck": "ck_bool_%(table_name)s_%(constraint_name)s"}``, then the + output of the following: + + op.add_column('t', 'x', Boolean(name='x')) + + will be:: + + CONSTRAINT ck_bool_t_x CHECK (x in (1, 0))) + + The function is rendered in the output of autogenerate when + a particular constraint name is already converted. + + """ + +def get_bind() -> "Connection": + """Return the current 'bind'. + + Under normal circumstances, this is the + :class:`~sqlalchemy.engine.Connection` currently being used + to emit SQL to the database. + + In a SQL script context, this value is ``None``. [TODO: verify this] + + """ + +def get_context(): + """Return the :class:`.MigrationContext` object that's + currently in use. + + """ + +def implementation_for(op_cls: Any) -> Callable: + """Register an implementation for a given :class:`.MigrateOperation`. + + This is part of the operation extensibility API. + + .. seealso:: + + :ref:`operation_plugins` - example of use + + """ + +def inline_literal( + value: Union[str, int], type_: None +) -> "_literal_bindparam": + """Produce an 'inline literal' expression, suitable for + using in an INSERT, UPDATE, or DELETE statement. + + When using Alembic in "offline" mode, CRUD operations + aren't compatible with SQLAlchemy's default behavior surrounding + literal values, + which is that they are converted into bound values and passed + separately into the ``execute()`` method of the DBAPI cursor. + An offline SQL + script needs to have these rendered inline. While it should + always be noted that inline literal values are an **enormous** + security hole in an application that handles untrusted input, + a schema migration is not run in this context, so + literals are safe to render inline, with the caveat that + advanced types like dates may not be supported directly + by SQLAlchemy. + + See :meth:`.execute` for an example usage of + :meth:`.inline_literal`. + + The environment can also be configured to attempt to render + "literal" values inline automatically, for those simple types + that are supported by the dialect; see + :paramref:`.EnvironmentContext.configure.literal_binds` for this + more recently added feature. + + :param value: The value to render. Strings, integers, and simple + numerics should be supported. Other types like boolean, + dates, etc. may or may not be supported yet by various + backends. + :param type\_: optional - a :class:`sqlalchemy.types.TypeEngine` + subclass stating the type of this value. In SQLAlchemy + expressions, this is usually derived automatically + from the Python type of the value itself, as well as + based on the context in which the value is used. + + .. seealso:: + + :paramref:`.EnvironmentContext.configure.literal_binds` + + """ + +def invoke(operation: "MigrateOperation") -> Any: + """Given a :class:`.MigrateOperation`, invoke it in terms of + this :class:`.Operations` instance. + + """ + +def register_operation(name: str, sourcename: Optional[str]) -> Callable: + """Register a new operation for this class. + + This method is normally used to add new operations + to the :class:`.Operations` class, and possibly the + :class:`.BatchOperations` class as well. All Alembic migration + operations are implemented via this system, however the system + is also available as a public API to facilitate adding custom + operations. + + .. seealso:: + + :ref:`operation_plugins` + + + """ + +def rename_table( + old_table_name: str, new_table_name: str, schema: Optional[str] +) -> Optional["Table"]: + """Emit an ALTER TABLE to rename a table. + + :param old_table_name: old name. + :param new_table_name: new name. + :param schema: Optional schema name to operate within. To control + quoting of the schema outside of the default behavior, use + the SQLAlchemy construct + :class:`~sqlalchemy.sql.elements.quoted_name`. + + """ diff --git a/alembic/operations/base.py b/alembic/operations/base.py index d4ec7b1..21f7a85 100644 --- a/alembic/operations/base.py +++ b/alembic/operations/base.py @@ -2,8 +2,12 @@ from contextlib import contextmanager import textwrap from typing import Any from typing import Callable +from typing import ForwardRef # noqa from typing import Iterator +from typing import List # noqa from typing import Optional +from typing import Sequence # noqa +from typing import Type # noqa from typing import TYPE_CHECKING from typing import Union @@ -14,9 +18,12 @@ from . import schemaobj from .. import util from ..util import sqla_compat from ..util.compat import inspect_formatargspec -from ..util.compat import inspect_getargspec +from ..util.compat import inspect_getfullargspec + +NoneType = type(None) if TYPE_CHECKING: + from sqlalchemy import Table # noqa from sqlalchemy.engine import Connection from .batch import BatchOperationsImpl @@ -107,7 +114,7 @@ class Operations(util.ModuleClsProxy): fn = getattr(op_cls, sourcename) source_name = fn.__name__ - spec = inspect_getargspec(fn) + spec = inspect_getfullargspec(fn) name_args = spec[0] assert name_args[0:2] == ["cls", "operations"] @@ -143,8 +150,10 @@ class Operations(util.ModuleClsProxy): "doc": fn.__doc__, } ) - globals_ = {"op_cls": op_cls} + globals_ = dict(globals()) + globals_.update({"op_cls": op_cls}) lcl = {} + exec(func_text, globals_, lcl) setattr(cls, name, lcl[name]) fn.__func__.__doc__ = ( diff --git a/alembic/py.typed b/alembic/py.typed new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/alembic/py.typed diff --git a/alembic/script/write_hooks.py b/alembic/script/write_hooks.py index 8f9e35e..8990148 100644 --- a/alembic/script/write_hooks.py +++ b/alembic/script/write_hooks.py @@ -32,6 +32,7 @@ def register(name: str) -> Callable: def decorate(fn): _registry[name] = fn + return fn return decorate @@ -112,7 +113,7 @@ def _parse_cmdline_options(cmdline_options_str: str, path: str) -> List[str]: @register("console_scripts") -def console_scripts(path, options): +def console_scripts(path, options, ignore_output=False): import pkg_resources try: @@ -128,6 +129,9 @@ def console_scripts(path, options): cmdline_options_str = options.get("options", "") cmdline_options_list = _parse_cmdline_options(cmdline_options_str, path) + kw = {} + if ignore_output: + kw["stdout"] = kw["stderr"] = subprocess.DEVNULL subprocess.run( [ sys.executable, @@ -137,4 +141,5 @@ def console_scripts(path, options): ] + cmdline_options_list, cwd=cwd, + **kw ) diff --git a/alembic/testing/__init__.py b/alembic/testing/__init__.py index 05b2e34..98ef82d 100644 --- a/alembic/testing/__init__.py +++ b/alembic/testing/__init__.py @@ -4,6 +4,7 @@ from sqlalchemy.testing import engines from sqlalchemy.testing import exclusions from sqlalchemy.testing import mock from sqlalchemy.testing import provide_metadata +from sqlalchemy.testing import skip_if from sqlalchemy.testing import uses_deprecated from sqlalchemy.testing.config import combinations from sqlalchemy.testing.config import fixture diff --git a/alembic/util/compat.py b/alembic/util/compat.py index a07813c..b87f8a6 100644 --- a/alembic/util/compat.py +++ b/alembic/util/compat.py @@ -1,48 +1,10 @@ -import collections -import inspect import io import os -from typing import Any -from typing import Callable -from typing import Dict -from typing import List -from typing import Optional -from typing import Type -is_posix = os.name == "posix" - -ArgSpec = collections.namedtuple( - "ArgSpec", ["args", "varargs", "keywords", "defaults"] -) - - -def inspect_getargspec(func: Callable) -> ArgSpec: - """getargspec based on fully vendored getfullargspec from Python 3.3.""" - - if inspect.ismethod(func): - func = func.__func__ # type: ignore - if not inspect.isfunction(func): - raise TypeError("{!r} is not a Python function".format(func)) +from sqlalchemy.util import inspect_getfullargspec # noqa +from sqlalchemy.util.compat import inspect_formatargspec # noqa - co = func.__code__ - if not inspect.iscode(co): - raise TypeError("{!r} is not a code object".format(co)) - - nargs = co.co_argcount - names = co.co_varnames - nkwargs = co.co_kwonlyargcount - args = list(names[:nargs]) - - nargs += nkwargs - varargs = None - if co.co_flags & inspect.CO_VARARGS: - varargs = co.co_varnames[nargs] - nargs = nargs + 1 - varkw = None - if co.co_flags & inspect.CO_VARKEYWORDS: - varkw = co.co_varnames[nargs] - - return ArgSpec(args, varargs, varkw, func.__defaults__) # type: ignore +is_posix = os.name == "posix" string_types = (str,) @@ -50,76 +12,6 @@ binary_type = bytes text_type = str -def _formatannotation(annotation, base_module=None): - """vendored from python 3.7""" - - if getattr(annotation, "__module__", None) == "typing": - return repr(annotation).replace("typing.", "") - if isinstance(annotation, type): - if annotation.__module__ in ("builtins", base_module): - return annotation.__qualname__ - return annotation.__module__ + "." + annotation.__qualname__ - return repr(annotation) - - -def inspect_formatargspec( - args: List[str], - varargs: Optional[str] = None, - varkw: Optional[str] = None, - defaults: Optional[Any] = None, - kwonlyargs: tuple = (), - kwonlydefaults: Dict[Any, Any] = {}, - annotations: Dict[Any, Any] = {}, - formatarg: Type[str] = str, - formatvarargs: Callable = lambda name: "*" + name, - formatvarkw: Callable = lambda name: "**" + name, - formatvalue: Callable = lambda value: "=" + repr(value), - formatreturns: Callable = lambda text: " -> " + text, - formatannotation: Callable = _formatannotation, -) -> str: - """Copy formatargspec from python 3.7 standard library. - - Python 3 has deprecated formatargspec and requested that Signature - be used instead, however this requires a full reimplementation - of formatargspec() in terms of creating Parameter objects and such. - Instead of introducing all the object-creation overhead and having - to reinvent from scratch, just copy their compatibility routine. - - """ - - def formatargandannotation(arg): - result = formatarg(arg) - if arg in annotations: - result += ": " + formatannotation(annotations[arg]) - return result - - specs = [] - if defaults: - firstdefault = len(args) - len(defaults) - for i, arg in enumerate(args): - spec = formatargandannotation(arg) - if defaults and i >= firstdefault: - spec = spec + formatvalue(defaults[i - firstdefault]) - specs.append(spec) - if varargs is not None: - specs.append(formatvarargs(formatargandannotation(varargs))) - else: - if kwonlyargs: - specs.append("*") - if kwonlyargs: - for kwonlyarg in kwonlyargs: - spec = formatargandannotation(kwonlyarg) - if kwonlydefaults and kwonlyarg in kwonlydefaults: - spec += formatvalue(kwonlydefaults[kwonlyarg]) - specs.append(spec) - if varkw is not None: - specs.append(formatvarkw(formatargandannotation(varkw))) - result = "(" + ", ".join(specs) + ")" - if "return" in annotations: - result += formatreturns(formatannotation(annotations["return"])) - return result - - # produce a wrapper that allows encoded text to stream # into a given buffer, but doesn't close it. # not sure of a more idiomatic approach to this. diff --git a/alembic/util/langhelpers.py b/alembic/util/langhelpers.py index 87a9aca..4db9a5f 100644 --- a/alembic/util/langhelpers.py +++ b/alembic/util/langhelpers.py @@ -20,7 +20,7 @@ from sqlalchemy.util import memoized_property # noqa from sqlalchemy.util import to_list # noqa from sqlalchemy.util import unique_list # noqa -from .compat import inspect_getargspec +from .compat import inspect_getfullargspec from .compat import string_types @@ -105,7 +105,7 @@ class ModuleClsProxy(metaclass=_ModuleClsMeta): translations = getattr(fn, "_legacy_translations", []) if translations: - spec = inspect_getargspec(fn) + spec = inspect_getfullargspec(fn) if spec[0] and spec[0][0] == "self": spec[0].pop(0) @@ -174,6 +174,7 @@ class ModuleClsProxy(metaclass=_ModuleClsMeta): } ) lcl = {} + exec(func_text, globals_, lcl) return lcl[name] diff --git a/alembic/util/sqla_compat.py b/alembic/util/sqla_compat.py index e1ccd41..b2b6f06 100644 --- a/alembic/util/sqla_compat.py +++ b/alembic/util/sqla_compat.py @@ -61,7 +61,7 @@ sqla_14 = _vers >= (1, 4) try: from sqlalchemy import Computed # noqa except ImportError: - Computed = None # type: ignore + Computed = type(None) # type: ignore has_computed = False has_computed_reflection = False else: @@ -71,7 +71,7 @@ else: try: from sqlalchemy import Identity # noqa except ImportError: - Identity = None # type: ignore + Identity = type(None) # type: ignore has_identity = False else: # attributes common to Indentity and Sequence diff --git a/docs/build/unreleased/py3_typing.rst b/docs/build/unreleased/py3_typing.rst index 7f8aa6c..fa4d2cf 100644 --- a/docs/build/unreleased/py3_typing.rst +++ b/docs/build/unreleased/py3_typing.rst @@ -1,8 +1,10 @@ .. change:: :tags: feature, general - pep-484 type annotations have been added throughout the library. This - should be helpful in providing Mypy and IDE support, however there is not - full support for Alembic's dynamically modified "op" namespace as of yet; a - future release will likely modify the approach used for importing this - namespace to be better compatible with pep-484 capabilities.
\ No newline at end of file + pep-484 type annotations have been added throughout the library. + Additionally, stub .pyi files have been added for the "dynamically" + generated Alembic modules ``alembic.op`` and ``alembic.config``, which + include complete function signatures and docstrings, so that the functions + in these namespaces will have both IDE support (vscode, pycharm, etc) as + well as support for typing tools like Mypy. The files themselves are + statically generated from their source functions within the source tree.
\ No newline at end of file @@ -47,6 +47,9 @@ install_requires = tz = python-dateutil +[options.package_data] +alembic = *.pyi, py.typed + [options.packages.find] exclude = test* diff --git a/tests/test_stubs.py b/tests/test_stubs.py new file mode 100644 index 0000000..c5186c3 --- /dev/null +++ b/tests/test_stubs.py @@ -0,0 +1,52 @@ +from pathlib import Path +import subprocess +import sys + +import alembic +from alembic.testing import eq_ +from alembic.testing import skip_if +from alembic.testing import TestBase + +_home = Path(__file__).parent.parent + + +def requirements(): + try: + import black # noqa + import zimports # noqa + + return False + except Exception: + return True + + +def run_command(file): + res = subprocess.run( + [ + sys.executable, + str((_home / "tools" / "write_pyi.py").relative_to(_home)), + "--stdout", + "--file", + file, + ], + stdout=subprocess.PIPE, + cwd=_home, + encoding="utf-8", + ) + return res + + +class TestStubFiles(TestBase): + @skip_if(requirements, "black and zimports are required for this test") + def test_op_pyi(self): + res = run_command("op") + generated = res.stdout + expected = Path(alembic.__file__).parent / "op.pyi" + eq_(generated, expected.read_text()) + + @skip_if(requirements, "black and zimports are required for this test") + def test_context_pyi(self): + res = run_command("context") + generated = res.stdout + expected = Path(alembic.__file__).parent / "context.pyi" + eq_(generated, expected.read_text()) diff --git a/tools/write_pyi.py b/tools/write_pyi.py new file mode 100644 index 0000000..6dfed09 --- /dev/null +++ b/tools/write_pyi.py @@ -0,0 +1,195 @@ +from argparse import ArgumentParser +from pathlib import Path +import re +import sys +from tempfile import NamedTemporaryFile +import textwrap + +from mako.pygen import PythonPrinter + +from alembic.operations.base import Operations +from alembic.runtime.environment import EnvironmentContext +from alembic.script.write_hooks import console_scripts +from alembic.util.compat import inspect_formatargspec +from alembic.util.compat import inspect_getfullargspec + +IGNORE_ITEMS = { + "op": {"context", "create_module_class_proxy"}, + "context": { + "config", + "create_module_class_proxy", + "get_impl", + "requires_connection", + "script", + }, +} + + +def generate_pyi_for_proxy( + cls: type, + progname: str, + source_path: Path, + destination_path: Path, + ignore_output: bool, + ignore_items: set, +): + + imports = [] + read_imports = False + with open(source_path) as read_file: + for line in read_file: + if line.startswith("# ### this file stubs are generated by"): + read_imports = True + elif line.startswith("### end imports ###"): + read_imports = False + break + elif read_imports: + imports.append(line.rstrip()) + + with open(destination_path, "w") as buf: + printer = PythonPrinter(buf) + + printer.writeline( + f"# ### this file stubs are generated by {progname} - do not edit ###" + ) + for line in imports: + buf.write(line + "\n") + printer.writeline("### end imports ###") + buf.write("\n\n") + + for name in dir(cls): + if name.startswith("_") or name in ignore_items: + continue + meth = getattr(cls, name) + if callable(meth): + _generate_stub_for_meth(cls, name, printer) + else: + _generate_stub_for_attr(cls, name, printer) + + printer.close() + + console_scripts( + str(destination_path), + {"entrypoint": "zimports", "options": "-e"}, + ignore_output=ignore_output, + ) + pyproject = Path(__file__).parent.parent / "pyproject.toml" + console_scripts( + str(destination_path), + {"entrypoint": "black", "options": f"--config {pyproject}"}, + ignore_output=ignore_output, + ) + + +def _generate_stub_for_attr(cls, name, printer): + printer.writeline(f"{name}: Any") + + +def _generate_stub_for_meth(cls, name, printer): + + fn = getattr(cls, name) + while hasattr(fn, "__wrapped__"): + fn = fn.__wrapped__ + + spec = inspect_getfullargspec(fn) + + name_args = spec[0] + assert name_args[0:1] == ["self"] or name_args[0:1] == ["cls"] + + name_args[0:1] = [] + + def _formatannotation(annotation, base_module=None): + if getattr(annotation, "__module__", None) == "typing": + retval = repr(annotation).replace("typing.", "") + elif isinstance(annotation, type): + if annotation.__module__ in ("builtins", base_module): + retval = annotation.__qualname__ + else: + retval = annotation.__module__ + "." + annotation.__qualname__ + else: + retval = repr(annotation) + + retval = re.sub( + r'ForwardRef\(([\'"].+?[\'"])\)', lambda m: m.group(1), retval + ) + retval = re.sub("NoneType", "None", retval) + return retval + + argspec = inspect_formatargspec( + *spec, formatvalue=lambda value: "", formatannotation=_formatannotation + ) + + func_text = textwrap.dedent( + """\ + def %(name)s%(argspec)s: + '''%(doc)s''' + """ + % { + "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 +): + progname = Path(sys.argv[0]).as_posix() + if not stdout: + generate_pyi_for_proxy( + cls_to_generate, + progname, + source_path=source_path, + destination_path=source_path, + ignore_output=False, + ignore_items=ignore_items, + ) + else: + with NamedTemporaryFile(delete=False, suffix=".pyi") as f: + f.close() + generate_pyi_for_proxy( + cls_to_generate, + progname, + source_path=source_path, + destination_path=f.name, + ignore_output=True, + ignore_items=ignore_items, + ) + f_path = Path(f.name) + 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"] + ) + if args.file in {"all", "context"}: + run_file( + location / "context.pyi", + EnvironmentContext, + args.stdout, + IGNORE_ITEMS["context"], + ) + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument( + "--file", + choices={"op", "context", "all"}, + default="all", + help="Which file to generate. Default is to regenerate all files", + ) + parser.add_argument( + "--stdout", + action="store_true", + help="Write to stdout instead of saving to file", + ) + args = parser.parse_args() + main(args) @@ -21,6 +21,8 @@ deps=pytest>4.6 sqlalchemy: sqlalchemy>=1.3.0 mako python-dateutil + zimports + black |