Add configurable revision_map_class option

alembic/script/__init__.py

@@ -1,4 +1,5 @@
 from .base import Script
 from .base import ScriptDirectory
+from .revision import RevisionMap
 
-__all__ = ["ScriptDirectory", "Script"]
+__all__ = ["ScriptDirectory", "Script", "RevisionMap"]

alembic/script/base.py

@@ -16,6 +16,7 @@
 from typing import Sequence
 from typing import Set
 from typing import Tuple
+from typing import Type
 from typing import TYPE_CHECKING
 from typing import Union
 
@@ -86,6 +87,9 @@ def __init__(
         messaging_opts: MessagingOptions = cast(
             "MessagingOptions", util.EMPTY_DICT
         ),
+        revision_map_class: Optional[
+            Type[revision.RevisionMap]
+        ] = None,
     ) -> None:
         self.dir = _preserving_path_as_str(dir)
         self.version_locations = [
@@ -95,7 +99,12 @@ def __init__(
         self.truncate_slug_length = truncate_slug_length or 40
         self.sourceless = sourceless
         self.output_encoding = output_encoding
-        self.revision_map = revision.RevisionMap(self._load_revisions)
+        revision_map_cls = (
+            revision_map_class
+            if revision_map_class is not None
+            else revision.RevisionMap
+        )
+        self.revision_map = revision_map_cls(self._load_revisions)
         self.timezone = timezone
         self.hooks = hooks
         self.recursive_version_locations = recursive_version_locations
@@ -183,6 +192,21 @@ def from_config(cls, config: Config) -> ScriptDirectory:
             sys.path[:0] = prepend_sys_path
 
         rvl = config.get_alembic_boolean_option("recursive_version_locations")
+
+        revision_map_class: Optional[Type[revision.RevisionMap]] = None
+        rmc = config.get_alembic_option("revision_map_class")
+        if rmc is not None:
+            resolved = util.resolve_dotted_name(rmc)
+            if not (
+                isinstance(resolved, type)
+                and issubclass(resolved, revision.RevisionMap)
+            ):
+                raise util.CommandError(
+                    f"revision_map_class {rmc!r} must be a subclass of "
+                    "alembic.script.revision.RevisionMap"
+                )
+            revision_map_class = resolved
+
         return ScriptDirectory(
             util.coerce_resource_to_filename(script_location),
             file_template=config.get_alembic_option(
@@ -198,6 +222,7 @@ def from_config(cls, config: Config) -> ScriptDirectory:
             hooks=config.get_hooks_list(),
             recursive_version_locations=rvl,
             messaging_opts=config.messaging_opts,
+            revision_map_class=revision_map_class,
         )
 
     @contextmanager

alembic/templates/async/alembic.ini.mako

@@ -79,6 +79,11 @@ path_separator = os
 # new in Alembic version 1.10
 # recursive_version_locations = false
 
+# specify a custom RevisionMap subclass for custom revision ordering logic.
+# the value is a dotted Python path in "module:ClassName" format.
+# the class must be a subclass of alembic.script.revision.RevisionMap.
+# revision_map_class = mypackage.custom:CustomRevisionMap
+
 # the output encoding used when revision files
 # are written from script.py.mako
 # output_encoding = utf-8

alembic/templates/generic/alembic.ini.mako

@@ -79,6 +79,11 @@ path_separator = os
 # new in Alembic version 1.10
 # recursive_version_locations = false
 
+# specify a custom RevisionMap subclass for custom revision ordering logic.
+# the value is a dotted Python path in "module:ClassName" format.
+# the class must be a subclass of alembic.script.revision.RevisionMap.
+# revision_map_class = mypackage.custom:CustomRevisionMap
+
 # the output encoding used when revision files
 # are written from script.py.mako
 # output_encoding = utf-8

alembic/templates/multidb/alembic.ini.mako

@@ -78,6 +78,11 @@ path_separator = os
 # new in Alembic version 1.10
 # recursive_version_locations = false
 
+# specify a custom RevisionMap subclass for custom revision ordering logic.
+# the value is a dotted Python path in "module:ClassName" format.
+# the class must be a subclass of alembic.script.revision.RevisionMap.
+# revision_map_class = mypackage.custom:CustomRevisionMap
+
 # the output encoding used when revision files
 # are written from script.py.mako
 # output_encoding = utf-8

alembic/templates/pyproject/pyproject.toml.mako

@@ -53,6 +53,11 @@ prepend_sys_path = [
 # new in Alembic version 1.10
 # recursive_version_locations = false
 
+# specify a custom RevisionMap subclass for custom revision ordering logic.
+# the value is a dotted Python path in "module:ClassName" format.
+# the class must be a subclass of alembic.script.revision.RevisionMap.
+# revision_map_class = "mypackage.custom:CustomRevisionMap"
+
 # the output encoding used when revision files
 # are written from script.py.mako
 # output_encoding = "utf-8"

alembic/templates/pyproject_async/pyproject.toml.mako

@@ -53,6 +53,11 @@ prepend_sys_path = [
 # new in Alembic version 1.10
 # recursive_version_locations = false
 
+# specify a custom RevisionMap subclass for custom revision ordering logic.
+# the value is a dotted Python path in "module:ClassName" format.
+# the class must be a subclass of alembic.script.revision.RevisionMap.
+# revision_map_class = "mypackage.custom:CustomRevisionMap"
+
 # the output encoding used when revision files
 # are written from script.py.mako
 # output_encoding = "utf-8"

alembic/util/__init__.py

@@ -29,5 +29,6 @@
 from .pyfiles import coerce_resource_to_filename as coerce_resource_to_filename
 from .pyfiles import load_python_file as load_python_file
 from .pyfiles import pyc_file_from_path as pyc_file_from_path
+from .pyfiles import resolve_dotted_name as resolve_dotted_name
 from .pyfiles import template_to_file as template_to_file
 from .sqla_compat import sqla_2 as sqla_2

alembic/util/pyfiles.py

@@ -21,6 +21,41 @@
 from .exc import CommandError
 
 
+def resolve_dotted_name(dotted_name: str) -> Any:
+    """Resolve a dotted name string to a Python object.
+
+    Accepts either ``"package.module:AttributeName"`` format (preferred,
+    following setuptools entry point conventions) or
+    ``"package.module.AttributeName"`` format.
+
+    """
+    if ":" in dotted_name:
+        module_part, _, attr_part = dotted_name.partition(":")
+    elif "." in dotted_name:
+        module_part, _, attr_part = dotted_name.rpartition(".")
+    else:
+        raise CommandError(
+            f"Could not resolve dotted name '{dotted_name}'; expected "
+            "format 'package.module:ClassName' or 'package.module.ClassName'"
+        )
+
+    try:
+        module = importlib.import_module(module_part)
+    except ImportError as ie:
+        raise CommandError(
+            f"Could not import module '{module_part}' "
+            f"from dotted name '{dotted_name}'"
+        ) from ie
+
+    try:
+        return getattr(module, attr_part)
+    except AttributeError as ae:
+        raise CommandError(
+            f"Module '{module_part}' has no attribute '{attr_part}' "
+            f"(from dotted name '{dotted_name}')"
+        ) from ae
+
+
 def template_to_file(
     template_file: Union[str, os.PathLike[str]],
     dest: Union[str, os.PathLike[str]],

docs/build/unreleased/revision_map_class.rst

@@ -0,0 +1,13 @@
+.. change::
+    :tags: usecase, scripts
+    :tickets:
+
+    Added a ``revision_map_class`` configuration option to
+    :class:`.ScriptDirectory`, allowing a custom subclass of
+    :class:`.RevisionMap` to be specified.  This enables external packages to
+    implement custom revision ordering logic, such as ordering based on git
+    history, without changes to Alembic core.  The option can be set in
+    ``alembic.ini`` or ``pyproject.toml`` using a dotted Python path in
+    ``"module:ClassName"`` format.  :class:`.RevisionMap` is now also exported
+    from the ``alembic.script`` package for convenient subclassing.
+

tests/test_config.py

@@ -46,6 +46,41 @@ def tearDown(self):
         clear_staging_env()
 
 
+class ResolveDottedNameTest(TestBase):
+    def test_colon_format(self):
+        result = util.resolve_dotted_name("collections:OrderedDict")
+        from collections import OrderedDict
+
+        assert result is OrderedDict
+
+    def test_dot_format(self):
+        result = util.resolve_dotted_name("collections.OrderedDict")
+        from collections import OrderedDict
+
+        assert result is OrderedDict
+
+    def test_invalid_no_separator(self):
+        with expect_raises_message(
+            util.CommandError,
+            r"Could not resolve dotted name",
+        ):
+            util.resolve_dotted_name("justaplainname")
+
+    def test_bad_module(self):
+        with expect_raises_message(
+            util.CommandError,
+            r"Could not import module",
+        ):
+            util.resolve_dotted_name("nonexistent.module:Foo")
+
+    def test_bad_attr(self):
+        with expect_raises_message(
+            util.CommandError,
+            r"has no attribute",
+        ):
+            util.resolve_dotted_name("collections:NoSuchThing")
+
+
 class ConfigTest(TestBase):
     def test_config_logging_with_file(self):
         buf = io.StringIO()
@@ -752,6 +787,83 @@ def test_setting(self):
         eq_(script.output_encoding, "latin-1")
 
 
+class RevisionMapClassTest(TestBase):
+    def setUp(self):
+        self.env = staging_env()
+        self.cfg = _no_sql_testing_config()
+
+    def tearDown(self):
+        clear_staging_env()
+
+    def test_default_revision_map_class(self):
+        from alembic.script.revision import RevisionMap
+
+        script = ScriptDirectory.from_config(self.cfg)
+        assert type(script.revision_map) is RevisionMap
+
+    def test_custom_class_via_constructor(self):
+        from alembic.script.revision import RevisionMap
+
+        class CustomRevisionMap(RevisionMap):
+            pass
+
+        script = ScriptDirectory(
+            self.cfg.get_main_option("script_location"),
+            revision_map_class=CustomRevisionMap,
+        )
+        assert type(script.revision_map) is CustomRevisionMap
+
+    def test_custom_class_from_config(self):
+        self.cfg.set_main_option(
+            "revision_map_class",
+            "alembic.script.revision:RevisionMap",
+        )
+        script = ScriptDirectory.from_config(self.cfg)
+        from alembic.script.revision import RevisionMap
+
+        assert type(script.revision_map) is RevisionMap
+
+    def test_invalid_not_a_subclass(self):
+        self.cfg.set_main_option(
+            "revision_map_class",
+            "collections:OrderedDict",
+        )
+        with expect_raises_message(
+            util.CommandError,
+            r"revision_map_class.*must be a subclass of "
+            r"alembic.script.revision.RevisionMap",
+        ):
+            ScriptDirectory.from_config(self.cfg)
+
+    def test_invalid_unresolvable_module(self):
+        self.cfg.set_main_option(
+            "revision_map_class",
+            "nonexistent.module:FakeClass",
+        )
+        with expect_raises_message(
+            util.CommandError,
+            r"Could not import module 'nonexistent.module'",
+        ):
+            ScriptDirectory.from_config(self.cfg)
+
+    def test_invalid_unresolvable_attr(self):
+        self.cfg.set_main_option(
+            "revision_map_class",
+            "alembic.script.revision:NoSuchClass",
+        )
+        with expect_raises_message(
+            util.CommandError,
+            r"has no attribute 'NoSuchClass'",
+        ):
+            ScriptDirectory.from_config(self.cfg)
+
+    def test_revision_map_exported_from_script_package(self):
+        from alembic.script import RevisionMap
+        from alembic.script.revision import RevisionMap as DirectRevisionMap
+
+        assert RevisionMap is DirectRevisionMap
+
+
 class CommandLineTest(TestBase):
     def test_register_command(self):
         cli = config.CommandLine()