Merge pull request #132 from schireson/dc/pyproject_config

fix: Handle pyproject.toml based alembic config.

Author: DanCardin

CHANGELOG.md

@@ -1,11 +1,18 @@
 # Changelog
 
-## [Unreleased](https://github.com/schireson/pytest-alembic/compare/v0.11.0...HEAD) (2025-05-16)
+### [v0.12.1](https://github.com/schireson/pytest-alembic/compare/v0.12.0...v0.12.1) (2025-05-22)
+
+#### Fixes
+
+* Handle pyproject.toml based alembic config.
+([82d9d62](https://github.com/schireson/pytest-alembic/commit/82d9d62024ccb86b655272a70e0d94374b5675a5))
+
+## [v0.12.0](https://github.com/schireson/pytest-alembic/compare/v0.11.1...v0.12.0) (2025-05-16)
 
 ### Fixes
 
 * Release 0.12.0
-([87dbf29](https://github.com/schireson/pytest-alembic/commit/87dbf29e24f47417a5821fb6908673058ead7d89))
+([4baee1e](https://github.com/schireson/pytest-alembic/commit/4baee1e85abe3733e6aa785c7704248f4fbb59b0))
 * Failing tests.
 ([1dc5b43](https://github.com/schireson/pytest-alembic/commit/1dc5b43d7e5b7530181808490cb31386631990a9))
 * Linting.
@@ -17,6 +24,13 @@
 * Updated pyproject.toml to be more flexible with poetry_core versioning.
 ([c7f25c3](https://github.com/schireson/pytest-alembic/commit/c7f25c39a795591a794f1ab89203f4a992012319))
 
+### [v0.11.1](https://github.com/schireson/pytest-alembic/compare/v0.11.0...v0.11.1) (2024-03-27)
+
+#### Fixes
+
+* Ensure branched revisions are upgraded individually once.
+([25b03a3](https://github.com/schireson/pytest-alembic/commit/25b03a3cac04258bbf3725f00e4794663808581a))
+
 ## [v0.11.0](https://github.com/schireson/pytest-alembic/compare/v0.10.7...v0.11.0) (2024-03-04)
 
 ### Fixes

pyproject.toml

@@ -1,41 +1,35 @@
 [tool.poetry]
 name = "pytest-alembic"
-version = "0.12.0"
+version = "0.12.1"
 description = "A pytest plugin for verifying alembic migrations."
-authors = [
-    "Dan Cardin <[email protected]>",
-]
+authors = ["Dan Cardin <[email protected]>"]
 license = "MIT"
-keywords = [ "pytest", "sqlalchemy", "alembic", "migration", "revision" ]
-classifiers = [ "Framework :: Pytest" ]
+keywords = ["pytest", "sqlalchemy", "alembic", "migration", "revision"]
+classifiers = ["Framework :: Pytest"]
 repository = "https://github.com/schireson/pytest-alembic"
-packages = [
-    { include = "pytest_alembic", from = "src" },
-]
+packages = [{ include = "pytest_alembic", from = "src" }]
 readme = 'README.md'
-include = [
-    "*.md",
-]
+include = ["*.md"]
 
 [tool.poetry.dependencies]
 python = ">=3.9, <4"
 
-pytest = {version = ">=7.0"}
+pytest = { version = ">=7.0" }
 alembic = "*"
 sqlalchemy = "*"
 
 [tool.poetry.group.dev.dependencies]
 asyncpg = "*"
-black = {version = "22.3.0", python = ">=3.6.2"}
-coverage = {version = ">=6.4.4", extras = ["toml"]}
+black = { version = "22.3.0", python = ">=3.6.2" }
+coverage = { version = ">=6.4.4", extras = ["toml"] }
 greenlet = ">=3.0"
-mypy = {version = ">=1.11", python = ">=3.8"}
+mypy = { version = ">=1.11", python = ">=3.8" }
 psycopg2-binary = "*"
-pytest = {version = ">=8.3.2", python = ">=3.8"}
+pytest = { version = ">=8.3.2", python = ">=3.8" }
 pytest-asyncio = "*"
-pytest-mock-resources = {version = ">=2.6.3", extras = ["docker"]}
-ruff = {version = '0.0.269'}
-sqlalchemy = {version = ">=1.4", extras = ["asyncio"]}
+pytest-mock-resources = { version = ">=2.6.3", extras = ["docker"] }
+ruff = { version = '0.0.269' }
+sqlalchemy = { version = ">=1.4", extras = ["asyncio"] }
 types-dataclasses = "^0.1.7"
 
 [tool.poetry.plugins.pytest11]
@@ -47,21 +41,54 @@ line_length = 100
 [tool.coverage.report]
 show_missing = true
 skip_covered = true
-exclude_lines = [
-  "pragma: no cover",
-  "if TYPE_CHECKING:",
-  "if __name__ == .__main__.:",
-]
+exclude_lines = ["pragma: no cover", "if TYPE_CHECKING:", "if __name__ == .__main__.:"]
 
 [tool.ruff]
 src = ["src", "tests"]
 target-version = "py37"
 line-length = 100
 select = [
-  "A", "B", "C", "D", "E", "F", "G", "I", "N", "Q", "S", "T", "W", "ARG", "BLE",
-  "DJ", "DTZ", "EM", "ERA", "EXE", "FBT", "ICN", "INP", "ISC", "NPY", "PD", "PGH",
-  "PIE", "PT", "PTH", "PYI", "RET", "RSE", "RUF", "SIM", "SLF", "TCH", "TID",
-  "TRY", "UP", "YTT"
+  "A",
+  "B",
+  "C",
+  "D",
+  "E",
+  "F",
+  "G",
+  "I",
+  "N",
+  "Q",
+  "S",
+  "T",
+  "W",
+  "ARG",
+  "BLE",
+  "DJ",
+  "DTZ",
+  "EM",
+  "ERA",
+  "EXE",
+  "FBT",
+  "ICN",
+  "INP",
+  "ISC",
+  "NPY",
+  "PD",
+  "PGH",
+  "PIE",
+  "PT",
+  "PTH",
+  "PYI",
+  "RET",
+  "RSE",
+  "RUF",
+  "SIM",
+  "SLF",
+  "TCH",
+  "TID",
+  "TRY",
+  "UP",
+  "YTT",
 ]
 ignore = ["E501", "S101", "D1", "TRY200", "B007", "B904"]
 
@@ -95,7 +122,7 @@ filterwarnings = [
   "ignore:unclosed connection.*:ResourceWarning",
   "ignore:pkg_resources is deprecated as an API:DeprecationWarning",
   "ignore:.*is deprecated and will be removed in Python 3.14.*:DeprecationWarning",
-  ]
+]
 
 [tool.mypy]
 strict_optional = true

src/pytest_alembic/config.py

@@ -2,9 +2,10 @@
 from typing import Any, cast, Dict, List, Optional, TYPE_CHECKING, Union
 
 import alembic.config
-from alembic.util import immutabledict
 
 if TYPE_CHECKING:
+    from alembic.util import immutabledict
+
     from pytest_alembic.revision_data import RevisionSpec
 
 
@@ -104,7 +105,7 @@ def from_raw_config(
         )
 
     def make_alembic_config(self, stdout):
-        file = (
+        ini_file = (
             self.config_options.get("file")
             or self.config_options.get("config_file_name")
             or "alembic.ini"
@@ -116,19 +117,24 @@ def make_alembic_config(self, stdout):
             alembic_config = self.alembic_config
             alembic_config.stdout = stdout
         else:
-            alembic_config = alembic.config.Config(file, stdout=stdout)
+            if _supports_toml():
+                alembic_config = alembic.config.Config(
+                    file_=ini_file, toml_file="pyproject.toml", stdout=stdout
+                )
+            else:
+                alembic_config = alembic.config.Config(ini_file, stdout=stdout)
 
         sqlalchemy_url = self.config_options.get("sqlalchemy.url")
         if sqlalchemy_url:
             alembic_config.set_main_option("sqlalchemy.url", sqlalchemy_url)
 
-        # Only set script_location if set.
         script_location = self.config_options.get("script_location")
-        if script_location:
-            alembic_config.set_section_option("alembic", "script_location", script_location)
-        elif not alembic_config.get_section_option("alembic", "script_location"):
+        if not script_location:
             # Or in the event that it's not set after already having loaded the config.
-            alembic_config.set_main_option("script_location", "migrations")
+            script_location = self._get_option(
+                alembic_config, "script_location", default="migrations"
+            )
+        alembic_config.set_main_option("script_location", script_location)
 
         target_metadata = self.config_options.get("target_metadata")
         alembic_config.attributes["target_metadata"] = target_metadata
@@ -141,6 +147,13 @@ def make_alembic_config(self, stdout):
 
         return alembic_config
 
+    @staticmethod
+    def _get_option(alembic_config: alembic.config.Config, key: str, *, default: str) -> str:
+        if _supports_toml():
+            get_alembic_option = getattr(alembic_config, "get_alembic_option")  # noqa: B009
+            return get_alembic_option(key, default)
+        return alembic_config.get_main_option(key, default)
+
 
 def duplicate_alembic_config(config: alembic.config.Config):
     return alembic.config.Config(
@@ -149,6 +162,10 @@ def duplicate_alembic_config(config: alembic.config.Config):
         output_buffer=config.output_buffer,
         stdout=config.stdout,
         cmd_opts=config.cmd_opts,
-        config_args=cast(immutabledict, config.config_args),
+        config_args=cast("immutabledict", config.config_args),
         attributes=config.attributes,
     )
+
+
+def _supports_toml() -> bool:
+    return hasattr(alembic.config.Config, "get_alembic_option")

tests/test_config.py

@@ -1,34 +1,37 @@
 from unittest import mock
 
+import alembic.config
+import pytest
+
 from pytest_alembic.config import Config
 
 
-def mock_open(*args, **kargs):
-    f_open = mock.mock_open(*args, **kargs)
-    f_open.return_value.__iter__ = lambda self: iter(self.readline, "")
-    return f_open
+def mock_patch(*contents):
+    m_open = mock.mock_open()
+    m_open.side_effect = [(mock.mock_open(read_data=string).return_value) for string in contents]
+    return mock.patch("builtins.open", m_open)
 
 
 def test_default_file():
     config = Config()
 
-    with mock.patch("builtins.open", mock_open(read_data="[alembic]")):
+    with mock_patch("[alembic]", b""):
         alembic_config = config.make_alembic_config(None)
     assert alembic_config.config_file_name == "alembic.ini"
 
 
 def test_set_file():
     config = Config(config_options={"file": "foo.ini"})
 
-    with mock.patch("builtins.open", mock_open(read_data="[alembic]")):
+    with mock_patch("[alembic]", b""):
         alembic_config = config.make_alembic_config(None)
     assert alembic_config.config_file_name == "foo.ini"
 
 
 def test_set_config_file_name():
     config = Config(config_options={"config_file_name": "foo.ini"})
 
-    with mock.patch("builtins.open", mock_open(read_data="[alembic]")):
+    with mock_patch("[alembic]", b""):
         alembic_config = config.make_alembic_config(None)
     assert alembic_config.config_file_name == "foo.ini"
 
@@ -43,3 +46,16 @@ def test_set_script_location():
     config = Config(config_options={"script_location": "alembic"})
     alembic_config = config.make_alembic_config(None)
     assert alembic_config.get_main_option("script_location") == "alembic"
+
+
+@pytest.mark.skipif(
+    not hasattr(alembic.config.Config, "get_alembic_option"),
+    reason="pyproject.toml is not supported in this version of alembic",
+)
+def test_pyproject_toml():
+    config = Config()
+
+    exists = mock.patch("alembic.config.Config._toml_file_path", return_value=True)
+    with exists, mock_patch("[alembic]", b"[tool.alembic]\nscript_location = 'asdf'"):
+        alembic_config = config.make_alembic_config(None)
+    assert alembic_config.get_main_option("script_location") == "asdf"