codekingpro/portable-devtools

1from __future__ import annotations2 3from argparse import ArgumentParser4from argparse import Namespace5from configparser import ConfigParser6import inspect7import os8import sys9from typing import Any10from typing import cast11from typing import Dict12from typing import Mapping13from typing import Optional14from typing import overload15from typing import Sequence16from typing import TextIO17from typing import Union18 19from typing_extensions import TypedDict20 21from . import __version__22from . import command23from . import util24from .util import compat25 26 27class Config:28    r"""Represent an Alembic configuration.29 30    Within an ``env.py`` script, this is available31    via the :attr:`.EnvironmentContext.config` attribute,32    which in turn is available at ``alembic.context``::33 34        from alembic import context35 36        some_param = context.config.get_main_option("my option")37 38    When invoking Alembic programmatically, a new39    :class:`.Config` can be created by passing40    the name of an .ini file to the constructor::41 42        from alembic.config import Config43        alembic_cfg = Config("/path/to/yourapp/alembic.ini")44 45    With a :class:`.Config` object, you can then46    run Alembic commands programmatically using the directives47    in :mod:`alembic.command`.48 49    The :class:`.Config` object can also be constructed without50    a filename.   Values can be set programmatically, and51    new sections will be created as needed::52 53        from alembic.config import Config54        alembic_cfg = Config()55        alembic_cfg.set_main_option("script_location", "myapp:migrations")56        alembic_cfg.set_main_option("sqlalchemy.url", "postgresql://foo/bar")57        alembic_cfg.set_section_option("mysection", "foo", "bar")58 59    .. warning::60 61       When using programmatic configuration, make sure the62       ``env.py`` file in use is compatible with the target configuration;63       including that the call to Python ``logging.fileConfig()`` is64       omitted if the programmatic configuration doesn't actually include65       logging directives.66 67    For passing non-string values to environments, such as connections and68    engines, use the :attr:`.Config.attributes` dictionary::69 70        with engine.begin() as connection:71            alembic_cfg.attributes['connection'] = connection72            command.upgrade(alembic_cfg, "head")73 74    :param file\_: name of the .ini file to open.75    :param ini_section: name of the main Alembic section within the76     .ini file77    :param output_buffer: optional file-like input buffer which78     will be passed to the :class:`.MigrationContext` - used to redirect79     the output of "offline generation" when using Alembic programmatically.80    :param stdout: buffer where the "print" output of commands will be sent.81     Defaults to ``sys.stdout``.82 83    :param config_args: A dictionary of keys and values that will be used84     for substitution in the alembic config file.  The dictionary as given85     is **copied** to a new one, stored locally as the attribute86     ``.config_args``. When the :attr:`.Config.file_config` attribute is87     first invoked, the replacement variable ``here`` will be added to this88     dictionary before the dictionary is passed to ``ConfigParser()``89     to parse the .ini file.90 91    :param attributes: optional dictionary of arbitrary Python keys/values,92     which will be populated into the :attr:`.Config.attributes` dictionary.93 94     .. seealso::95 96        :ref:`connection_sharing`97 98    """99 100    def __init__(101        self,102        file_: Union[str, os.PathLike[str], None] = None,103        ini_section: str = "alembic",104        output_buffer: Optional[TextIO] = None,105        stdout: TextIO = sys.stdout,106        cmd_opts: Optional[Namespace] = None,107        config_args: Mapping[str, Any] = util.immutabledict(),108        attributes: Optional[Dict[str, Any]] = None,109    ) -> None:110        """Construct a new :class:`.Config`"""111        self.config_file_name = file_112        self.config_ini_section = ini_section113        self.output_buffer = output_buffer114        self.stdout = stdout115        self.cmd_opts = cmd_opts116        self.config_args = dict(config_args)117        if attributes:118            self.attributes.update(attributes)119 120    cmd_opts: Optional[Namespace] = None121    """The command-line options passed to the ``alembic`` script.122 123    Within an ``env.py`` script this can be accessed via the124    :attr:`.EnvironmentContext.config` attribute.125 126    .. seealso::127 128        :meth:`.EnvironmentContext.get_x_argument`129 130    """131 132    config_file_name: Union[str, os.PathLike[str], None] = None133    """Filesystem path to the .ini file in use."""134 135    config_ini_section: str = None  # type:ignore[assignment]136    """Name of the config file section to read basic configuration137    from.  Defaults to ``alembic``, that is the ``[alembic]`` section138    of the .ini file.  This value is modified using the ``-n/--name``139    option to the Alembic runner.140 141    """142 143    @util.memoized_property144    def attributes(self) -> Dict[str, Any]:145        """A Python dictionary for storage of additional state.146 147 148        This is a utility dictionary which can include not just strings but149        engines, connections, schema objects, or anything else.150        Use this to pass objects into an env.py script, such as passing151        a :class:`sqlalchemy.engine.base.Connection` when calling152        commands from :mod:`alembic.command` programmatically.153 154        .. seealso::155 156            :ref:`connection_sharing`157 158            :paramref:`.Config.attributes`159 160        """161        return {}162 163    def print_stdout(self, text: str, *arg: Any) -> None:164        """Render a message to standard out.165 166        When :meth:`.Config.print_stdout` is called with additional args167        those arguments will formatted against the provided text,168        otherwise we simply output the provided text verbatim.169 170        This is a no-op when the``quiet`` messaging option is enabled.171 172        e.g.::173 174            >>> config.print_stdout('Some text %s', 'arg')175            Some Text arg176 177        """178 179        if arg:180            output = str(text) % arg181        else:182            output = str(text)183 184        util.write_outstream(self.stdout, output, "\n", **self.messaging_opts)185 186    @util.memoized_property187    def file_config(self) -> ConfigParser:188        """Return the underlying ``ConfigParser`` object.189 190        Direct access to the .ini file is available here,191        though the :meth:`.Config.get_section` and192        :meth:`.Config.get_main_option`193        methods provide a possibly simpler interface.194 195        """196 197        if self.config_file_name:198            here = os.path.abspath(os.path.dirname(self.config_file_name))199        else:200            here = ""201        self.config_args["here"] = here202        file_config = ConfigParser(self.config_args)203        if self.config_file_name:204            compat.read_config_parser(file_config, [self.config_file_name])205        else:206            file_config.add_section(self.config_ini_section)207        return file_config208 209    def get_template_directory(self) -> str:210        """Return the directory where Alembic setup templates are found.211 212        This method is used by the alembic ``init`` and ``list_templates``213        commands.214 215        """216        import alembic217 218        package_dir = os.path.abspath(os.path.dirname(alembic.__file__))219        return os.path.join(package_dir, "templates")220 221    @overload222    def get_section(223        self, name: str, default: None = ...224    ) -> Optional[Dict[str, str]]:225        ...226 227    # "default" here could also be a TypeVar228    # _MT = TypeVar("_MT", bound=Mapping[str, str]),229    # however mypy wasn't handling that correctly (pyright was)230    @overload231    def get_section(232        self, name: str, default: Dict[str, str]233    ) -> Dict[str, str]:234        ...235 236    @overload237    def get_section(238        self, name: str, default: Mapping[str, str]239    ) -> Union[Dict[str, str], Mapping[str, str]]:240        ...241 242    def get_section(243        self, name: str, default: Optional[Mapping[str, str]] = None244    ) -> Optional[Mapping[str, str]]:245        """Return all the configuration options from a given .ini file section246        as a dictionary.247 248        If the given section does not exist, the value of ``default``249        is returned, which is expected to be a dictionary or other mapping.250 251        """252        if not self.file_config.has_section(name):253            return default254 255        return dict(self.file_config.items(name))256 257    def set_main_option(self, name: str, value: str) -> None:258        """Set an option programmatically within the 'main' section.259 260        This overrides whatever was in the .ini file.261 262        :param name: name of the value263 264        :param value: the value.  Note that this value is passed to265         ``ConfigParser.set``, which supports variable interpolation using266         pyformat (e.g. ``%(some_value)s``).   A raw percent sign not part of267         an interpolation symbol must therefore be escaped, e.g. ``%%``.268         The given value may refer to another value already in the file269         using the interpolation format.270 271        """272        self.set_section_option(self.config_ini_section, name, value)273 274    def remove_main_option(self, name: str) -> None:275        self.file_config.remove_option(self.config_ini_section, name)276 277    def set_section_option(self, section: str, name: str, value: str) -> None:278        """Set an option programmatically within the given section.279 280        The section is created if it doesn't exist already.281        The value here will override whatever was in the .ini282        file.283 284        :param section: name of the section285 286        :param name: name of the value287 288        :param value: the value.  Note that this value is passed to289         ``ConfigParser.set``, which supports variable interpolation using290         pyformat (e.g. ``%(some_value)s``).   A raw percent sign not part of291         an interpolation symbol must therefore be escaped, e.g. ``%%``.292         The given value may refer to another value already in the file293         using the interpolation format.294 295        """296 297        if not self.file_config.has_section(section):298            self.file_config.add_section(section)299        self.file_config.set(section, name, value)300 301    def get_section_option(302        self, section: str, name: str, default: Optional[str] = None303    ) -> Optional[str]:304        """Return an option from the given section of the .ini file."""305        if not self.file_config.has_section(section):306            raise util.CommandError(307                "No config file %r found, or file has no "308                "'[%s]' section" % (self.config_file_name, section)309            )310        if self.file_config.has_option(section, name):311            return self.file_config.get(section, name)312        else:313            return default314 315    @overload316    def get_main_option(self, name: str, default: str) -> str:317        ...318 319    @overload320    def get_main_option(321        self, name: str, default: Optional[str] = None322    ) -> Optional[str]:323        ...324 325    def get_main_option(326        self, name: str, default: Optional[str] = None327    ) -> Optional[str]:328        """Return an option from the 'main' section of the .ini file.329 330        This defaults to being a key from the ``[alembic]``331        section, unless the ``-n/--name`` flag were used to332        indicate a different section.333 334        """335        return self.get_section_option(self.config_ini_section, name, default)336 337    @util.memoized_property338    def messaging_opts(self) -> MessagingOptions:339        """The messaging options."""340        return cast(341            MessagingOptions,342            util.immutabledict(343                {"quiet": getattr(self.cmd_opts, "quiet", False)}344            ),345        )346 347 348class MessagingOptions(TypedDict, total=False):349    quiet: bool350 351 352class CommandLine:353    def __init__(self, prog: Optional[str] = None) -> None:354        self._generate_args(prog)355 356    def _generate_args(self, prog: Optional[str]) -> None:357        def add_options(358            fn: Any, parser: Any, positional: Any, kwargs: Any359        ) -> None:360            kwargs_opts = {361                "template": (362                    "-t",363                    "--template",364                    dict(365                        default="generic",366                        type=str,367                        help="Setup template for use with 'init'",368                    ),369                ),370                "message": (371                    "-m",372                    "--message",373                    dict(374                        type=str, help="Message string to use with 'revision'"375                    ),376                ),377                "sql": (378                    "--sql",379                    dict(380                        action="store_true",381                        help="Don't emit SQL to database - dump to "382                        "standard output/file instead. See docs on "383                        "offline mode.",384                    ),385                ),386                "tag": (387                    "--tag",388                    dict(389                        type=str,390                        help="Arbitrary 'tag' name - can be used by "391                        "custom env.py scripts.",392                    ),393                ),394                "head": (395                    "--head",396                    dict(397                        type=str,398                        help="Specify head revision or <branchname>@head "399                        "to base new revision on.",400                    ),401                ),402                "splice": (403                    "--splice",404                    dict(405                        action="store_true",406                        help="Allow a non-head revision as the "407                        "'head' to splice onto",408                    ),409                ),410                "depends_on": (411                    "--depends-on",412                    dict(413                        action="append",414                        help="Specify one or more revision identifiers "415                        "which this revision should depend on.",416                    ),417                ),418                "rev_id": (419                    "--rev-id",420                    dict(421                        type=str,422                        help="Specify a hardcoded revision id instead of "423                        "generating one",424                    ),425                ),426                "version_path": (427                    "--version-path",428                    dict(429                        type=str,430                        help="Specify specific path from config for "431                        "version file",432                    ),433                ),434                "branch_label": (435                    "--branch-label",436                    dict(437                        type=str,438                        help="Specify a branch label to apply to the "439                        "new revision",440                    ),441                ),442                "verbose": (443                    "-v",444                    "--verbose",445                    dict(action="store_true", help="Use more verbose output"),446                ),447                "resolve_dependencies": (448                    "--resolve-dependencies",449                    dict(450                        action="store_true",451                        help="Treat dependency versions as down revisions",452                    ),453                ),454                "autogenerate": (455                    "--autogenerate",456                    dict(457                        action="store_true",458                        help="Populate revision script with candidate "459                        "migration operations, based on comparison "460                        "of database to model.",461                    ),462                ),463                "rev_range": (464                    "-r",465                    "--rev-range",466                    dict(467                        action="store",468                        help="Specify a revision range; "469                        "format is [start]:[end]",470                    ),471                ),472                "indicate_current": (473                    "-i",474                    "--indicate-current",475                    dict(476                        action="store_true",477                        help="Indicate the current revision",478                    ),479                ),480                "purge": (481                    "--purge",482                    dict(483                        action="store_true",484                        help="Unconditionally erase the version table "485                        "before stamping",486                    ),487                ),488                "package": (489                    "--package",490                    dict(491                        action="store_true",492                        help="Write empty __init__.py files to the "493                        "environment and version locations",494                    ),495                ),496            }497            positional_help = {498                "directory": "location of scripts directory",499                "revision": "revision identifier",500                "revisions": "one or more revisions, or 'heads' for all heads",501            }502            for arg in kwargs:503                if arg in kwargs_opts:504                    args = kwargs_opts[arg]505                    args, kw = args[0:-1], args[-1]506                    parser.add_argument(*args, **kw)507 508            for arg in positional:509                if (510                    arg == "revisions"511                    or fn in positional_translations512                    and positional_translations[fn][arg] == "revisions"513                ):514                    subparser.add_argument(515                        "revisions",516                        nargs="+",517                        help=positional_help.get("revisions"),518                    )519                else:520                    subparser.add_argument(arg, help=positional_help.get(arg))521 522        parser = ArgumentParser(prog=prog)523 524        parser.add_argument(525            "--version", action="version", version="%%(prog)s %s" % __version__526        )527        parser.add_argument(528            "-c",529            "--config",530            type=str,531            default=os.environ.get("ALEMBIC_CONFIG", "alembic.ini"),532            help="Alternate config file; defaults to value of "533            'ALEMBIC_CONFIG environment variable, or "alembic.ini"',534        )535        parser.add_argument(536            "-n",537            "--name",538            type=str,539            default="alembic",540            help="Name of section in .ini file to " "use for Alembic config",541        )542        parser.add_argument(543            "-x",544            action="append",545            help="Additional arguments consumed by "546            "custom env.py scripts, e.g. -x "547            "setting1=somesetting -x setting2=somesetting",548        )549        parser.add_argument(550            "--raiseerr",551            action="store_true",552            help="Raise a full stack trace on error",553        )554        parser.add_argument(555            "-q",556            "--quiet",557            action="store_true",558            help="Do not log to std output.",559        )560        subparsers = parser.add_subparsers()561 562        positional_translations: Dict[Any, Any] = {563            command.stamp: {"revision": "revisions"}564        }565 566        for fn in [getattr(command, n) for n in dir(command)]:567            if (568                inspect.isfunction(fn)569                and fn.__name__[0] != "_"570                and fn.__module__ == "alembic.command"571            ):572                spec = compat.inspect_getfullargspec(fn)573                if spec[3] is not None:574                    positional = spec[0][1 : -len(spec[3])]575                    kwarg = spec[0][-len(spec[3]) :]576                else:577                    positional = spec[0][1:]578                    kwarg = []579 580                if fn in positional_translations:581                    positional = [582                        positional_translations[fn].get(name, name)583                        for name in positional584                    ]585 586                # parse first line(s) of helptext without a line break587                help_ = fn.__doc__588                if help_:589                    help_text = []590                    for line in help_.split("\n"):591                        if not line.strip():592                            break593                        else:594                            help_text.append(line.strip())595                else:596                    help_text = []597                subparser = subparsers.add_parser(598                    fn.__name__, help=" ".join(help_text)599                )600                add_options(fn, subparser, positional, kwarg)601                subparser.set_defaults(cmd=(fn, positional, kwarg))602        self.parser = parser603 604    def run_cmd(self, config: Config, options: Namespace) -> None:605        fn, positional, kwarg = options.cmd606 607        try:608            fn(609                config,610                *[getattr(options, k, None) for k in positional],611                **{k: getattr(options, k, None) for k in kwarg},612            )613        except util.CommandError as e:614            if options.raiseerr:615                raise616            else:617                util.err(str(e), **config.messaging_opts)618 619    def main(self, argv: Optional[Sequence[str]] = None) -> None:620        options = self.parser.parse_args(argv)621        if not hasattr(options, "cmd"):622            # see http://bugs.python.org/issue9253, argparse623            # behavior changed incompatibly in py3.3624            self.parser.error("too few arguments")625        else:626            cfg = Config(627                file_=options.config,628                ini_section=options.name,629                cmd_opts=options,630            )631            self.run_cmd(cfg, options)632 633 634def main(635    argv: Optional[Sequence[str]] = None,636    prog: Optional[str] = None,637    **kwargs: Any,638) -> None:639    """The console runner function for Alembic."""640 641    CommandLine(prog=prog).main(argv=argv)642 643 644if __name__ == "__main__":645    main()646