from __future__ import annotations from importlib import metadata import logging import re from re import Pattern from types import ModuleType from typing import Callable from typing import TYPE_CHECKING from .. import util from ..util import DispatchPriority from ..util import PriorityDispatcher if TYPE_CHECKING: from ..util import PriorityDispatchResult _all_plugins: dict[str, Plugin] = {} log = logging.getLogger(__name__) class Plugin: """Describe a series of functions that are pulled in as a plugin. This is initially to provide for portable lists of autogenerate comparison functions, however the setup for a plugin can run any other kinds of global registration as well. .. versionadded:: 1.18.0 """ def __init__(self, name: str): self.name = name log.info("setup plugin %s", name) if name in _all_plugins: raise ValueError(f"A plugin named {name} is already registered") _all_plugins[name] = self self.autogenerate_comparators = PriorityDispatcher() def remove(self) -> None: """remove this plugin""" del _all_plugins[self.name] def add_autogenerate_comparator( self, fn: Callable[..., PriorityDispatchResult], compare_target: str, compare_element: str | None = None, *, qualifier: str = "default", priority: DispatchPriority = DispatchPriority.MEDIUM, ) -> None: """Register an autogenerate comparison function. See the section :ref:`plugins_registering_autogenerate` for detailed examples on how to use this method. :param fn: The comparison function to register. The function receives arguments specific to the type of comparison being performed and should return a :class:`.PriorityDispatchResult` value. :param compare_target: The type of comparison being performed (e.g., ``"table"``, ``"column"``, ``"type"``). :param compare_element: Optional sub-element being compared within the target type. :param qualifier: Database dialect qualifier. Use ``"default"`` for all dialects, or specify a dialect name like ``"postgresql"`` to register a dialect-specific handler. Defaults to ``"default"``. :param priority: Execution priority for this comparison function. Functions are executed in priority order from :attr:`.DispatchPriority.FIRST` to :attr:`.DispatchPriority.LAST`. Defaults to :attr:`.DispatchPriority.MEDIUM`. """ self.autogenerate_comparators.dispatch_for( compare_target, subgroup=compare_element, priority=priority, qualifier=qualifier, )(fn) @classmethod def populate_autogenerate_priority_dispatch( cls, comparators: PriorityDispatcher, include_plugins: list[str] ) -> None: """Populate all current autogenerate comparison functions into a given PriorityDispatcher.""" exclude: set[Pattern[str]] = set() include: dict[str, Pattern[str]] = {} matched_expressions: set[str] = set() for name in include_plugins: # the expression as written is retained as the dictionary key # so that error messages report what was actually passed if name.startswith("~"): exclude.add(_make_re(_translate_legacy_name(name[1:]))) else: include[name] = _make_re(_translate_legacy_name(name)) for plugin in _all_plugins.values(): if any(excl.match(plugin.name) for excl in exclude): continue include_matches = [ incl for incl in include if include[incl].match(plugin.name) ] if not include_matches: continue else: matched_expressions.update(include_matches) log.info("setting up autogenerate plugin %s", plugin.name) comparators.populate_with(plugin.autogenerate_comparators) never_matched = set(include).difference(matched_expressions) if never_matched: raise util.CommandError( f"Did not locate plugins: {', '.join(never_matched)}" ) @classmethod def setup_plugin_from_module(cls, module: ModuleType, name: str) -> None: """Call the ``setup()`` function of a plugin module, identified by passing the module object itself. E.g.:: from alembic.runtime.plugins import Plugin import myproject.alembic_plugin # Register the plugin manually Plugin.setup_plugin_from_module( myproject.alembic_plugin, "myproject.custom_operations" ) This will generate a new :class:`.Plugin` object with the given name, which will register itself in the global list of plugins. Then the module's ``setup()`` function is invoked, passing that :class:`.Plugin` object. This exact process is invoked automatically at import time for any plugin module that is published via the ``alembic.plugins`` entrypoint. """ module.setup(Plugin(name)) _legacy_plugin_names = { # renamed in 1.19.2 when the plugin became opt-in; see #1859. the # "alembic.ext." namespace is deliberately outside of the # "alembic.autogenerate.*" wildcard. "alembic.autogenerate.checkconstraint_byname": ( "alembic.ext.checkconstraint_byname" ), } def _translate_legacy_name(name: str) -> str: """Resolve a plugin name that has been renamed to its current name. This allows an ``env.py`` written against an earlier Alembic release to keep working, whether the name is used to include or to exclude the plugin. Wildcard expressions are never translated, as only fully stated names are eligible. """ return _legacy_plugin_names.get(name, name) def _make_re(name: str) -> Pattern[str]: tokens = name.split(".") reg = r"" for token in tokens: if token == "*": reg += r"\..+?" elif token.isidentifier(): reg += r"\." + token else: raise ValueError(f"Invalid plugin expression {name!r}") # omit leading r'\.' return re.compile(f"^{reg[2:]}$") def _setup() -> None: # setup third party plugins for entrypoint in metadata.entry_points(group="alembic.plugins"): for mod in entrypoint.load(): Plugin.setup_plugin_from_module(mod, entrypoint.name) _setup()