config.py 33 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026
  1. from __future__ import annotations
  2. from argparse import ArgumentParser
  3. from argparse import Namespace
  4. from configparser import ConfigParser
  5. import inspect
  6. import os
  7. from pathlib import Path
  8. import re
  9. import sys
  10. from typing import Any
  11. from typing import cast
  12. from typing import Dict
  13. from typing import Mapping
  14. from typing import Optional
  15. from typing import overload
  16. from typing import Protocol
  17. from typing import Sequence
  18. from typing import TextIO
  19. from typing import Union
  20. from typing_extensions import TypedDict
  21. from . import __version__
  22. from . import command
  23. from . import util
  24. from .util import compat
  25. from .util.pyfiles import _preserving_path_as_str
  26. class Config:
  27. r"""Represent an Alembic configuration.
  28. Within an ``env.py`` script, this is available
  29. via the :attr:`.EnvironmentContext.config` attribute,
  30. which in turn is available at ``alembic.context``::
  31. from alembic import context
  32. some_param = context.config.get_main_option("my option")
  33. When invoking Alembic programmatically, a new
  34. :class:`.Config` can be created by passing
  35. the name of an .ini file to the constructor::
  36. from alembic.config import Config
  37. alembic_cfg = Config("/path/to/yourapp/alembic.ini")
  38. With a :class:`.Config` object, you can then
  39. run Alembic commands programmatically using the directives
  40. in :mod:`alembic.command`.
  41. The :class:`.Config` object can also be constructed without
  42. a filename. Values can be set programmatically, and
  43. new sections will be created as needed::
  44. from alembic.config import Config
  45. alembic_cfg = Config()
  46. alembic_cfg.set_main_option("script_location", "myapp:migrations")
  47. alembic_cfg.set_main_option("sqlalchemy.url", "postgresql://foo/bar")
  48. alembic_cfg.set_section_option("mysection", "foo", "bar")
  49. .. warning::
  50. When using programmatic configuration, make sure the
  51. ``env.py`` file in use is compatible with the target configuration;
  52. including that the call to Python ``logging.fileConfig()`` is
  53. omitted if the programmatic configuration doesn't actually include
  54. logging directives.
  55. For passing non-string values to environments, such as connections and
  56. engines, use the :attr:`.Config.attributes` dictionary::
  57. with engine.begin() as connection:
  58. alembic_cfg.attributes['connection'] = connection
  59. command.upgrade(alembic_cfg, "head")
  60. :param file\_: name of the .ini file to open if an ``alembic.ini`` is
  61. to be used. This should refer to the ``alembic.ini`` file, either as
  62. a filename or a full path to the file. This filename if passed must refer
  63. to an **ini file in ConfigParser format** only.
  64. :param toml\_file: name of the pyproject.toml file to open if a
  65. ``pyproject.toml`` file is to be used. This should refer to the
  66. ``pyproject.toml`` file, either as a filename or a full path to the file.
  67. This file must be in toml format. Both :paramref:`.Config.file\_` and
  68. :paramref:`.Config.toml\_file` may be passed simultaneously, or
  69. exclusively.
  70. .. versionadded:: 1.16.0
  71. :param ini_section: name of the main Alembic section within the
  72. .ini file
  73. :param output_buffer: optional file-like input buffer which
  74. will be passed to the :class:`.MigrationContext` - used to redirect
  75. the output of "offline generation" when using Alembic programmatically.
  76. :param stdout: buffer where the "print" output of commands will be sent.
  77. Defaults to ``sys.stdout``.
  78. :param config_args: A dictionary of keys and values that will be used
  79. for substitution in the alembic config file, as well as the pyproject.toml
  80. file, depending on which / both are used. The dictionary as given is
  81. **copied** to two new, independent dictionaries, stored locally under the
  82. attributes ``.config_args`` and ``.toml_args``. Both of these
  83. dictionaries will also be populated with the replacement variable
  84. ``%(here)s``, which refers to the location of the .ini and/or .toml file
  85. as appropriate.
  86. :param attributes: optional dictionary of arbitrary Python keys/values,
  87. which will be populated into the :attr:`.Config.attributes` dictionary.
  88. .. seealso::
  89. :ref:`connection_sharing`
  90. """
  91. def __init__(
  92. self,
  93. file_: Union[str, os.PathLike[str], None] = None,
  94. toml_file: Union[str, os.PathLike[str], None] = None,
  95. ini_section: str = "alembic",
  96. output_buffer: Optional[TextIO] = None,
  97. stdout: TextIO = sys.stdout,
  98. cmd_opts: Optional[Namespace] = None,
  99. config_args: Mapping[str, Any] = util.immutabledict(),
  100. attributes: Optional[Dict[str, Any]] = None,
  101. ) -> None:
  102. """Construct a new :class:`.Config`"""
  103. self.config_file_name = (
  104. _preserving_path_as_str(file_) if file_ else None
  105. )
  106. self.toml_file_name = (
  107. _preserving_path_as_str(toml_file) if toml_file else None
  108. )
  109. self.config_ini_section = ini_section
  110. self.output_buffer = output_buffer
  111. self.stdout = stdout
  112. self.cmd_opts = cmd_opts
  113. self.config_args = dict(config_args)
  114. self.toml_args = dict(config_args)
  115. if attributes:
  116. self.attributes.update(attributes)
  117. cmd_opts: Optional[Namespace] = None
  118. """The command-line options passed to the ``alembic`` script.
  119. Within an ``env.py`` script this can be accessed via the
  120. :attr:`.EnvironmentContext.config` attribute.
  121. .. seealso::
  122. :meth:`.EnvironmentContext.get_x_argument`
  123. """
  124. config_file_name: Optional[str] = None
  125. """Filesystem path to the .ini file in use."""
  126. toml_file_name: Optional[str] = None
  127. """Filesystem path to the pyproject.toml file in use.
  128. .. versionadded:: 1.16.0
  129. """
  130. @property
  131. def _config_file_path(self) -> Optional[Path]:
  132. if self.config_file_name is None:
  133. return None
  134. return Path(self.config_file_name)
  135. @property
  136. def _toml_file_path(self) -> Optional[Path]:
  137. if self.toml_file_name is None:
  138. return None
  139. return Path(self.toml_file_name)
  140. config_ini_section: str = None # type:ignore[assignment]
  141. """Name of the config file section to read basic configuration
  142. from. Defaults to ``alembic``, that is the ``[alembic]`` section
  143. of the .ini file. This value is modified using the ``-n/--name``
  144. option to the Alembic runner.
  145. """
  146. @util.memoized_property
  147. def attributes(self) -> Dict[str, Any]:
  148. """A Python dictionary for storage of additional state.
  149. This is a utility dictionary which can include not just strings but
  150. engines, connections, schema objects, or anything else.
  151. Use this to pass objects into an env.py script, such as passing
  152. a :class:`sqlalchemy.engine.base.Connection` when calling
  153. commands from :mod:`alembic.command` programmatically.
  154. .. seealso::
  155. :ref:`connection_sharing`
  156. :paramref:`.Config.attributes`
  157. """
  158. return {}
  159. def print_stdout(self, text: str, *arg: Any) -> None:
  160. """Render a message to standard out.
  161. When :meth:`.Config.print_stdout` is called with additional args
  162. those arguments will formatted against the provided text,
  163. otherwise we simply output the provided text verbatim.
  164. This is a no-op when the``quiet`` messaging option is enabled.
  165. e.g.::
  166. >>> config.print_stdout('Some text %s', 'arg')
  167. Some Text arg
  168. """
  169. if arg:
  170. output = str(text) % arg
  171. else:
  172. output = str(text)
  173. util.write_outstream(self.stdout, output, "\n", **self.messaging_opts)
  174. @util.memoized_property
  175. def file_config(self) -> ConfigParser:
  176. """Return the underlying ``ConfigParser`` object.
  177. Dir*-ect access to the .ini file is available here,
  178. though the :meth:`.Config.get_section` and
  179. :meth:`.Config.get_main_option`
  180. methods provide a possibly simpler interface.
  181. """
  182. if self._config_file_path:
  183. here = self._config_file_path.absolute().parent
  184. else:
  185. here = Path()
  186. self.config_args["here"] = here.as_posix()
  187. file_config = ConfigParser(self.config_args)
  188. if self._config_file_path:
  189. compat.read_config_parser(file_config, [self._config_file_path])
  190. else:
  191. file_config.add_section(self.config_ini_section)
  192. return file_config
  193. @util.memoized_property
  194. def toml_alembic_config(self) -> Mapping[str, Any]:
  195. """Return a dictionary of the [tool.alembic] section from
  196. pyproject.toml"""
  197. if self._toml_file_path and self._toml_file_path.exists():
  198. here = self._toml_file_path.absolute().parent
  199. self.toml_args["here"] = here.as_posix()
  200. with open(self._toml_file_path, "rb") as f:
  201. toml_data = compat.tomllib.load(f)
  202. data = toml_data.get("tool", {}).get("alembic", {})
  203. if not isinstance(data, dict):
  204. raise util.CommandError("Incorrect TOML format")
  205. return data
  206. else:
  207. return {}
  208. def get_template_directory(self) -> str:
  209. """Return the directory where Alembic setup templates are found.
  210. This method is used by the alembic ``init`` and ``list_templates``
  211. commands.
  212. """
  213. import alembic
  214. package_dir = Path(alembic.__file__).absolute().parent
  215. return str(package_dir / "templates")
  216. def _get_template_path(self) -> Path:
  217. """Return the directory where Alembic setup templates are found.
  218. This method is used by the alembic ``init`` and ``list_templates``
  219. commands.
  220. .. versionadded:: 1.16.0
  221. """
  222. return Path(self.get_template_directory())
  223. @overload
  224. def get_section(
  225. self, name: str, default: None = ...
  226. ) -> Optional[Dict[str, str]]: ...
  227. # "default" here could also be a TypeVar
  228. # _MT = TypeVar("_MT", bound=Mapping[str, str]),
  229. # however mypy wasn't handling that correctly (pyright was)
  230. @overload
  231. def get_section(
  232. self, name: str, default: Dict[str, str]
  233. ) -> Dict[str, str]: ...
  234. @overload
  235. def get_section(
  236. self, name: str, default: Mapping[str, str]
  237. ) -> Union[Dict[str, str], Mapping[str, str]]: ...
  238. def get_section(
  239. self, name: str, default: Optional[Mapping[str, str]] = None
  240. ) -> Optional[Mapping[str, str]]:
  241. """Return all the configuration options from a given .ini file section
  242. as a dictionary.
  243. If the given section does not exist, the value of ``default``
  244. is returned, which is expected to be a dictionary or other mapping.
  245. """
  246. if not self.file_config.has_section(name):
  247. return default
  248. return dict(self.file_config.items(name))
  249. def set_main_option(self, name: str, value: str) -> None:
  250. """Set an option programmatically within the 'main' section.
  251. This overrides whatever was in the .ini file.
  252. :param name: name of the value
  253. :param value: the value. Note that this value is passed to
  254. ``ConfigParser.set``, which supports variable interpolation using
  255. pyformat (e.g. ``%(some_value)s``). A raw percent sign not part of
  256. an interpolation symbol must therefore be escaped, e.g. ``%%``.
  257. The given value may refer to another value already in the file
  258. using the interpolation format.
  259. """
  260. self.set_section_option(self.config_ini_section, name, value)
  261. def remove_main_option(self, name: str) -> None:
  262. self.file_config.remove_option(self.config_ini_section, name)
  263. def set_section_option(self, section: str, name: str, value: str) -> None:
  264. """Set an option programmatically within the given section.
  265. The section is created if it doesn't exist already.
  266. The value here will override whatever was in the .ini
  267. file.
  268. Does **NOT** consume from the pyproject.toml file.
  269. .. seealso::
  270. :meth:`.Config.get_alembic_option` - includes pyproject support
  271. :param section: name of the section
  272. :param name: name of the value
  273. :param value: the value. Note that this value is passed to
  274. ``ConfigParser.set``, which supports variable interpolation using
  275. pyformat (e.g. ``%(some_value)s``). A raw percent sign not part of
  276. an interpolation symbol must therefore be escaped, e.g. ``%%``.
  277. The given value may refer to another value already in the file
  278. using the interpolation format.
  279. """
  280. if not self.file_config.has_section(section):
  281. self.file_config.add_section(section)
  282. self.file_config.set(section, name, value)
  283. def get_section_option(
  284. self, section: str, name: str, default: Optional[str] = None
  285. ) -> Optional[str]:
  286. """Return an option from the given section of the .ini file."""
  287. if not self.file_config.has_section(section):
  288. raise util.CommandError(
  289. "No config file %r found, or file has no "
  290. "'[%s]' section" % (self.config_file_name, section)
  291. )
  292. if self.file_config.has_option(section, name):
  293. return self.file_config.get(section, name)
  294. else:
  295. return default
  296. @overload
  297. def get_main_option(self, name: str, default: str) -> str: ...
  298. @overload
  299. def get_main_option(
  300. self, name: str, default: Optional[str] = None
  301. ) -> Optional[str]: ...
  302. def get_main_option(
  303. self, name: str, default: Optional[str] = None
  304. ) -> Optional[str]:
  305. """Return an option from the 'main' section of the .ini file.
  306. This defaults to being a key from the ``[alembic]``
  307. section, unless the ``-n/--name`` flag were used to
  308. indicate a different section.
  309. Does **NOT** consume from the pyproject.toml file.
  310. .. seealso::
  311. :meth:`.Config.get_alembic_option` - includes pyproject support
  312. """
  313. return self.get_section_option(self.config_ini_section, name, default)
  314. @overload
  315. def get_alembic_option(self, name: str, default: str) -> str: ...
  316. @overload
  317. def get_alembic_option(
  318. self, name: str, default: Optional[str] = None
  319. ) -> Optional[str]: ...
  320. def get_alembic_option(
  321. self, name: str, default: Optional[str] = None
  322. ) -> Union[
  323. None, str, list[str], dict[str, str], list[dict[str, str]], int
  324. ]:
  325. """Return an option from the "[alembic]" or "[tool.alembic]" section
  326. of the configparser-parsed .ini file (e.g. ``alembic.ini``) or
  327. toml-parsed ``pyproject.toml`` file.
  328. The value returned is expected to be None, string, list of strings,
  329. or dictionary of strings. Within each type of string value, the
  330. ``%(here)s`` token is substituted out with the absolute path of the
  331. ``pyproject.toml`` file, as are other tokens which are extracted from
  332. the :paramref:`.Config.config_args` dictionary.
  333. Searches always prioritize the configparser namespace first, before
  334. searching in the toml namespace.
  335. If Alembic was run using the ``-n/--name`` flag to indicate an
  336. alternate main section name, this is taken into account **only** for
  337. the configparser-parsed .ini file. The section name in toml is always
  338. ``[tool.alembic]``.
  339. .. versionadded:: 1.16.0
  340. """
  341. if self.file_config.has_option(self.config_ini_section, name):
  342. return self.file_config.get(self.config_ini_section, name)
  343. else:
  344. return self._get_toml_config_value(name, default=default)
  345. def get_alembic_boolean_option(self, name: str) -> bool:
  346. if self.file_config.has_option(self.config_ini_section, name):
  347. return (
  348. self.file_config.get(self.config_ini_section, name) == "true"
  349. )
  350. else:
  351. value = self.toml_alembic_config.get(name, False)
  352. if not isinstance(value, bool):
  353. raise util.CommandError(
  354. f"boolean value expected for TOML parameter {name!r}"
  355. )
  356. return value
  357. def _get_toml_config_value(
  358. self, name: str, default: Optional[Any] = None
  359. ) -> Union[
  360. None, str, list[str], dict[str, str], list[dict[str, str]], int
  361. ]:
  362. USE_DEFAULT = object()
  363. value: Union[None, str, list[str], dict[str, str], int] = (
  364. self.toml_alembic_config.get(name, USE_DEFAULT)
  365. )
  366. if value is USE_DEFAULT:
  367. return default
  368. if value is not None:
  369. if isinstance(value, str):
  370. value = value % (self.toml_args)
  371. elif isinstance(value, list):
  372. if value and isinstance(value[0], dict):
  373. value = [
  374. {k: v % (self.toml_args) for k, v in dv.items()}
  375. for dv in value
  376. ]
  377. else:
  378. value = cast(
  379. "list[str]", [v % (self.toml_args) for v in value]
  380. )
  381. elif isinstance(value, dict):
  382. value = cast(
  383. "dict[str, str]",
  384. {k: v % (self.toml_args) for k, v in value.items()},
  385. )
  386. elif isinstance(value, int):
  387. return value
  388. else:
  389. raise util.CommandError(
  390. f"unsupported TOML value type for key: {name!r}"
  391. )
  392. return value
  393. @util.memoized_property
  394. def messaging_opts(self) -> MessagingOptions:
  395. """The messaging options."""
  396. return cast(
  397. MessagingOptions,
  398. util.immutabledict(
  399. {"quiet": getattr(self.cmd_opts, "quiet", False)}
  400. ),
  401. )
  402. def _get_file_separator_char(self, *names: str) -> Optional[str]:
  403. for name in names:
  404. separator = self.get_main_option(name)
  405. if separator is not None:
  406. break
  407. else:
  408. return None
  409. split_on_path = {
  410. "space": " ",
  411. "newline": "\n",
  412. "os": os.pathsep,
  413. ":": ":",
  414. ";": ";",
  415. }
  416. try:
  417. sep = split_on_path[separator]
  418. except KeyError as ke:
  419. raise ValueError(
  420. "'%s' is not a valid value for %s; "
  421. "expected 'space', 'newline', 'os', ':', ';'"
  422. % (separator, name)
  423. ) from ke
  424. else:
  425. if name == "version_path_separator":
  426. util.warn_deprecated(
  427. "The version_path_separator configuration parameter "
  428. "is deprecated; please use path_separator"
  429. )
  430. return sep
  431. def get_version_locations_list(self) -> Optional[list[str]]:
  432. version_locations_str = self.file_config.get(
  433. self.config_ini_section, "version_locations", fallback=None
  434. )
  435. if version_locations_str:
  436. split_char = self._get_file_separator_char(
  437. "path_separator", "version_path_separator"
  438. )
  439. if split_char is None:
  440. # legacy behaviour for backwards compatibility
  441. util.warn_deprecated(
  442. "No path_separator found in configuration; "
  443. "falling back to legacy splitting on spaces/commas "
  444. "for version_locations. Consider adding "
  445. "path_separator=os to Alembic config."
  446. )
  447. _split_on_space_comma = re.compile(r", *|(?: +)")
  448. return _split_on_space_comma.split(version_locations_str)
  449. else:
  450. return [
  451. x.strip()
  452. for x in version_locations_str.split(split_char)
  453. if x
  454. ]
  455. else:
  456. return cast(
  457. "list[str]",
  458. self._get_toml_config_value("version_locations", None),
  459. )
  460. def get_prepend_sys_paths_list(self) -> Optional[list[str]]:
  461. prepend_sys_path_str = self.file_config.get(
  462. self.config_ini_section, "prepend_sys_path", fallback=None
  463. )
  464. if prepend_sys_path_str:
  465. split_char = self._get_file_separator_char("path_separator")
  466. if split_char is None:
  467. # legacy behaviour for backwards compatibility
  468. util.warn_deprecated(
  469. "No path_separator found in configuration; "
  470. "falling back to legacy splitting on spaces, commas, "
  471. "and colons for prepend_sys_path. Consider adding "
  472. "path_separator=os to Alembic config."
  473. )
  474. _split_on_space_comma_colon = re.compile(r", *|(?: +)|\:")
  475. return _split_on_space_comma_colon.split(prepend_sys_path_str)
  476. else:
  477. return [
  478. x.strip()
  479. for x in prepend_sys_path_str.split(split_char)
  480. if x
  481. ]
  482. else:
  483. return cast(
  484. "list[str]",
  485. self._get_toml_config_value("prepend_sys_path", None),
  486. )
  487. def get_hooks_list(self) -> list[PostWriteHookConfig]:
  488. hooks: list[PostWriteHookConfig] = []
  489. if not self.file_config.has_section("post_write_hooks"):
  490. toml_hook_config = cast(
  491. "list[dict[str, str]]",
  492. self._get_toml_config_value("post_write_hooks", []),
  493. )
  494. for cfg in toml_hook_config:
  495. opts = dict(cfg)
  496. opts["_hook_name"] = opts.pop("name")
  497. hooks.append(opts)
  498. else:
  499. _split_on_space_comma = re.compile(r", *|(?: +)")
  500. ini_hook_config = self.get_section("post_write_hooks", {})
  501. names = _split_on_space_comma.split(
  502. ini_hook_config.get("hooks", "")
  503. )
  504. for name in names:
  505. if not name:
  506. continue
  507. opts = {
  508. key[len(name) + 1 :]: ini_hook_config[key]
  509. for key in ini_hook_config
  510. if key.startswith(name + ".")
  511. }
  512. opts["_hook_name"] = name
  513. hooks.append(opts)
  514. return hooks
  515. PostWriteHookConfig = Mapping[str, str]
  516. class MessagingOptions(TypedDict, total=False):
  517. quiet: bool
  518. class CommandFunction(Protocol):
  519. """A function that may be registered in the CLI as an alembic command.
  520. It must be a named function and it must accept a :class:`.Config` object
  521. as the first argument.
  522. .. versionadded:: 1.15.3
  523. """
  524. __name__: str
  525. def __call__(self, config: Config, *args: Any, **kwargs: Any) -> Any: ...
  526. class CommandLine:
  527. """Provides the command line interface to Alembic."""
  528. def __init__(self, prog: Optional[str] = None) -> None:
  529. self._generate_args(prog)
  530. _KWARGS_OPTS = {
  531. "template": (
  532. "-t",
  533. "--template",
  534. dict(
  535. default="generic",
  536. type=str,
  537. help="Setup template for use with 'init'",
  538. ),
  539. ),
  540. "message": (
  541. "-m",
  542. "--message",
  543. dict(type=str, help="Message string to use with 'revision'"),
  544. ),
  545. "sql": (
  546. "--sql",
  547. dict(
  548. action="store_true",
  549. help="Don't emit SQL to database - dump to "
  550. "standard output/file instead. See docs on "
  551. "offline mode.",
  552. ),
  553. ),
  554. "tag": (
  555. "--tag",
  556. dict(
  557. type=str,
  558. help="Arbitrary 'tag' name - can be used by "
  559. "custom env.py scripts.",
  560. ),
  561. ),
  562. "head": (
  563. "--head",
  564. dict(
  565. type=str,
  566. help="Specify head revision or <branchname>@head "
  567. "to base new revision on.",
  568. ),
  569. ),
  570. "splice": (
  571. "--splice",
  572. dict(
  573. action="store_true",
  574. help="Allow a non-head revision as the 'head' to splice onto",
  575. ),
  576. ),
  577. "depends_on": (
  578. "--depends-on",
  579. dict(
  580. action="append",
  581. help="Specify one or more revision identifiers "
  582. "which this revision should depend on.",
  583. ),
  584. ),
  585. "rev_id": (
  586. "--rev-id",
  587. dict(
  588. type=str,
  589. help="Specify a hardcoded revision id instead of "
  590. "generating one",
  591. ),
  592. ),
  593. "version_path": (
  594. "--version-path",
  595. dict(
  596. type=str,
  597. help="Specify specific path from config for version file",
  598. ),
  599. ),
  600. "branch_label": (
  601. "--branch-label",
  602. dict(
  603. type=str,
  604. help="Specify a branch label to apply to the new revision",
  605. ),
  606. ),
  607. "verbose": (
  608. "-v",
  609. "--verbose",
  610. dict(action="store_true", help="Use more verbose output"),
  611. ),
  612. "resolve_dependencies": (
  613. "--resolve-dependencies",
  614. dict(
  615. action="store_true",
  616. help="Treat dependency versions as down revisions",
  617. ),
  618. ),
  619. "autogenerate": (
  620. "--autogenerate",
  621. dict(
  622. action="store_true",
  623. help="Populate revision script with candidate "
  624. "migration operations, based on comparison "
  625. "of database to model.",
  626. ),
  627. ),
  628. "rev_range": (
  629. "-r",
  630. "--rev-range",
  631. dict(
  632. action="store",
  633. help="Specify a revision range; format is [start]:[end]",
  634. ),
  635. ),
  636. "indicate_current": (
  637. "-i",
  638. "--indicate-current",
  639. dict(
  640. action="store_true",
  641. help="Indicate the current revision",
  642. ),
  643. ),
  644. "purge": (
  645. "--purge",
  646. dict(
  647. action="store_true",
  648. help="Unconditionally erase the version table before stamping",
  649. ),
  650. ),
  651. "package": (
  652. "--package",
  653. dict(
  654. action="store_true",
  655. help="Write empty __init__.py files to the "
  656. "environment and version locations",
  657. ),
  658. ),
  659. }
  660. _POSITIONAL_OPTS = {
  661. "directory": dict(help="location of scripts directory"),
  662. "revision": dict(
  663. help="revision identifier",
  664. ),
  665. "revisions": dict(
  666. nargs="+",
  667. help="one or more revisions, or 'heads' for all heads",
  668. ),
  669. }
  670. _POSITIONAL_TRANSLATIONS: dict[Any, dict[str, str]] = {
  671. command.stamp: {"revision": "revisions"}
  672. }
  673. def _generate_args(self, prog: Optional[str]) -> None:
  674. parser = ArgumentParser(prog=prog)
  675. parser.add_argument(
  676. "--version", action="version", version="%%(prog)s %s" % __version__
  677. )
  678. parser.add_argument(
  679. "-c",
  680. "--config",
  681. action="append",
  682. help="Alternate config file; defaults to value of "
  683. 'ALEMBIC_CONFIG environment variable, or "alembic.ini". '
  684. "May also refer to pyproject.toml file. May be specified twice "
  685. "to reference both files separately",
  686. )
  687. parser.add_argument(
  688. "-n",
  689. "--name",
  690. type=str,
  691. default="alembic",
  692. help="Name of section in .ini file to use for Alembic config "
  693. "(only applies to configparser config, not toml)",
  694. )
  695. parser.add_argument(
  696. "-x",
  697. action="append",
  698. help="Additional arguments consumed by "
  699. "custom env.py scripts, e.g. -x "
  700. "setting1=somesetting -x setting2=somesetting",
  701. )
  702. parser.add_argument(
  703. "--raiseerr",
  704. action="store_true",
  705. help="Raise a full stack trace on error",
  706. )
  707. parser.add_argument(
  708. "-q",
  709. "--quiet",
  710. action="store_true",
  711. help="Do not log to std output.",
  712. )
  713. self.subparsers = parser.add_subparsers()
  714. alembic_commands = (
  715. cast(CommandFunction, fn)
  716. for fn in (getattr(command, name) for name in dir(command))
  717. if (
  718. inspect.isfunction(fn)
  719. and fn.__name__[0] != "_"
  720. and fn.__module__ == "alembic.command"
  721. )
  722. )
  723. for fn in alembic_commands:
  724. self.register_command(fn)
  725. self.parser = parser
  726. def register_command(self, fn: CommandFunction) -> None:
  727. """Registers a function as a CLI subcommand. The subcommand name
  728. matches the function name, the arguments are extracted from the
  729. signature and the help text is read from the docstring.
  730. .. versionadded:: 1.15.3
  731. .. seealso::
  732. :ref:`custom_commandline`
  733. """
  734. positional, kwarg, help_text = self._inspect_function(fn)
  735. subparser = self.subparsers.add_parser(fn.__name__, help=help_text)
  736. subparser.set_defaults(cmd=(fn, positional, kwarg))
  737. for arg in kwarg:
  738. if arg in self._KWARGS_OPTS:
  739. kwarg_opt = self._KWARGS_OPTS[arg]
  740. args, opts = kwarg_opt[0:-1], kwarg_opt[-1]
  741. subparser.add_argument(*args, **opts) # type:ignore
  742. for arg in positional:
  743. opts = self._POSITIONAL_OPTS.get(arg, {})
  744. subparser.add_argument(arg, **opts) # type:ignore
  745. def _inspect_function(self, fn: CommandFunction) -> tuple[Any, Any, str]:
  746. spec = compat.inspect_getfullargspec(fn)
  747. if spec[3] is not None:
  748. positional = spec[0][1 : -len(spec[3])]
  749. kwarg = spec[0][-len(spec[3]) :]
  750. else:
  751. positional = spec[0][1:]
  752. kwarg = []
  753. if fn in self._POSITIONAL_TRANSLATIONS:
  754. positional = [
  755. self._POSITIONAL_TRANSLATIONS[fn].get(name, name)
  756. for name in positional
  757. ]
  758. # parse first line(s) of helptext without a line break
  759. help_ = fn.__doc__
  760. if help_:
  761. help_lines = []
  762. for line in help_.split("\n"):
  763. if not line.strip():
  764. break
  765. else:
  766. help_lines.append(line.strip())
  767. else:
  768. help_lines = []
  769. help_text = " ".join(help_lines)
  770. return positional, kwarg, help_text
  771. def run_cmd(self, config: Config, options: Namespace) -> None:
  772. fn, positional, kwarg = options.cmd
  773. try:
  774. fn(
  775. config,
  776. *[getattr(options, k, None) for k in positional],
  777. **{k: getattr(options, k, None) for k in kwarg},
  778. )
  779. except util.CommandError as e:
  780. if options.raiseerr:
  781. raise
  782. else:
  783. util.err(str(e), **config.messaging_opts)
  784. def _inis_from_config(self, options: Namespace) -> tuple[str, str]:
  785. names = options.config
  786. alembic_config_env = os.environ.get("ALEMBIC_CONFIG")
  787. if (
  788. alembic_config_env
  789. and os.path.basename(alembic_config_env) == "pyproject.toml"
  790. ):
  791. default_pyproject_toml = alembic_config_env
  792. default_alembic_config = "alembic.ini"
  793. elif alembic_config_env:
  794. default_pyproject_toml = "pyproject.toml"
  795. default_alembic_config = alembic_config_env
  796. else:
  797. default_alembic_config = "alembic.ini"
  798. default_pyproject_toml = "pyproject.toml"
  799. if not names:
  800. return default_pyproject_toml, default_alembic_config
  801. toml = ini = None
  802. for name in names:
  803. if os.path.basename(name) == "pyproject.toml":
  804. if toml is not None:
  805. raise util.CommandError(
  806. "pyproject.toml indicated more than once"
  807. )
  808. toml = name
  809. else:
  810. if ini is not None:
  811. raise util.CommandError(
  812. "only one ini file may be indicated"
  813. )
  814. ini = name
  815. return toml if toml else default_pyproject_toml, (
  816. ini if ini else default_alembic_config
  817. )
  818. def main(self, argv: Optional[Sequence[str]] = None) -> None:
  819. """Executes the command line with the provided arguments."""
  820. options = self.parser.parse_args(argv)
  821. if not hasattr(options, "cmd"):
  822. # see http://bugs.python.org/issue9253, argparse
  823. # behavior changed incompatibly in py3.3
  824. self.parser.error("too few arguments")
  825. else:
  826. toml, ini = self._inis_from_config(options)
  827. cfg = Config(
  828. file_=ini,
  829. toml_file=toml,
  830. ini_section=options.name,
  831. cmd_opts=options,
  832. )
  833. self.run_cmd(cfg, options)
  834. def main(
  835. argv: Optional[Sequence[str]] = None,
  836. prog: Optional[str] = None,
  837. **kwargs: Any,
  838. ) -> None:
  839. """The console runner function for Alembic."""
  840. CommandLine(prog=prog).main(argv=argv)
  841. if __name__ == "__main__":
  842. main()