FEAT: automatically update settings in pyproject.toml (#143)

* FEAT: automatically update `black` configuration
* FEAT: automatically update `nbqa.addopts` settings
* MAINT: switch to `tomlkit`
  Just like `ruamel.yaml`, `tomlkit` perserves formatting of the TOML file when dumping it

.cspell.json

@@ -111,6 +111,7 @@
     "showcode",
     "showtags",
     "taplo",
+    "tomlkit",
     "unittests",
     "venv"
   ]

pyproject.toml

@@ -25,12 +25,12 @@ exclude = '''
 include = '\.pyi?$'
 preview = true
 target-version = [
-    'py310',
-    'py311',
-    'py36',
-    'py37',
-    'py38',
-    'py39',
+    "py310",
+    "py311",
+    "py36",
+    "py37",
+    "py38",
+    "py39",
 ]
 
 [tool.isort]

setup.cfg

@@ -36,7 +36,7 @@ install_requires =
     pip-tools
     PyYAML
     ruamel.yaml  # better YAML dumping
-    toml
+    tomlkit
 packages = find:
 package_dir =
     =src

src/repoma/check_dev_files/black.py

@@ -1,91 +1,76 @@
-"""Check :file:`pyproject.toml` black config."""
-from textwrap import dedent
-from typing import List, Optional
+"""Update :file:`pyproject.toml` black configuration."""
 
 from ruamel.yaml.comments import CommentedMap
 
 from repoma.errors import PrecommitError
-from repoma.utilities import CONFIG_PATH, natural_sorting
+from repoma.utilities import CONFIG_PATH
 from repoma.utilities.executor import Executor
 from repoma.utilities.precommit import (
     find_repo,
     load_round_trip_precommit_config,
     update_precommit_hook,
     update_single_hook_precommit_repo,
 )
-from repoma.utilities.pyproject import load_pyproject
+from repoma.utilities.pyproject import (
+    get_sub_table,
+    load_pyproject,
+    to_toml_array,
+    write_pyproject,
+)
 from repoma.utilities.setup_cfg import get_supported_python_versions
 
 
 def main() -> None:
     if not CONFIG_PATH.pyproject.exists():
         return
-    config = _load_black_config()
     executor = Executor()
-    executor(_check_line_length, config)
-    executor(_check_activate_preview, config)
-    executor(_check_option_ordering, config)
-    executor(_check_target_versions, config)
-    executor(_check_pyproject)
+    executor(_remove_outdated_settings)
+    executor(_update_black_settings)
+    executor(_update_nbqa_settings)
     executor(_update_precommit_repo)
     executor(_update_precommit_nbqa_hook)
     executor.finalize()
 
 
-def _load_black_config(content: Optional[str] = None) -> dict:
-    config = load_pyproject(content)
-    return config.get("tool", {}).get("black", {})
-
-
-def _check_activate_preview(config: dict) -> None:
-    expected_option = "preview"
-    if config.get(expected_option) is not True:
-        raise PrecommitError(dedent(f"""
-            An option in pyproject.toml is wrong or missing. Should be:
-
-            [tool.black]
-            {expected_option} = true
-            """).strip())
-
-
-def _check_line_length(config: dict) -> None:
-    if config.get("line-length") is not None:
-        raise PrecommitError(
-            "pyproject.toml should not specify a line-length (default to 88)."
+def _remove_outdated_settings() -> None:
+    pyproject = load_pyproject()
+    settings = get_sub_table(pyproject, "tool.black", create=True)
+    forbidden_options = ("line-length",)
+    removed_options = set()
+    for option in forbidden_options:
+        if option in settings:
+            removed_options.add(option)
+            settings.remove(option)
+    if removed_options:
+        write_pyproject(pyproject)
+        msg = (
+            f"Removed {', '.join(sorted(removed_options))} option from black"
+            f" configuration in {CONFIG_PATH.pyproject}"
         )
+        raise PrecommitError(msg)
 
 
-def _check_option_ordering(config: dict) -> None:
-    options = list(config)
-    sorted_options = sorted(config, key=natural_sorting)
-    if sorted_options != options:
-        error_message = dedent("""
-            Options in pyproject.toml should be alphabetically sorted:
-
-            [tool.black]
-            """).strip()
-        for option in sorted_options:
-            error_message += f"\n{option} = ..."
-        raise PrecommitError(error_message)
+def _update_black_settings() -> None:
+    pyproject = load_pyproject()
+    settings = get_sub_table(pyproject, "tool.black", create=True)
+    versions = get_supported_python_versions()
+    target_version = to_toml_array(sorted("py" + v.replace(".", "") for v in versions))
+    minimal_settings = {
+        "preview": True,
+        "target-version": target_version,
+    }
+    if not __complies(settings, minimal_settings):
+        settings.update(minimal_settings)
+        write_pyproject(pyproject)
+        msg = f"Updated black configuration in {CONFIG_PATH.pyproject}"
+        raise PrecommitError(msg)
 
 
-def _check_target_versions(config: dict) -> None:
-    target_versions = config.get("target-version", [])
-    supported_python_versions = get_supported_python_versions()
-    expected_target_versions = sorted(
-        "py" + s.replace(".", "") for s in supported_python_versions
-    )
-    if target_versions != expected_target_versions:
-        error_message = dedent("""
-            Black target versions in pyproject.toml should be as follows:
-
-            [tool.black]
-            target-version = [
-            """).strip()
-        for version in expected_target_versions:
-            error_message += f"\n    '{version}',"
-        error_message += "\n]"
-        raise PrecommitError(error_message)
+def __complies(settings: dict, minimal_settings: dict) -> bool:
+    for key, value in minimal_settings.items():
+        if settings.get(key) != value:
+            return False
+    return True
 
 
 def _update_precommit_repo() -> None:
@@ -106,29 +91,25 @@ def _update_precommit_nbqa_hook() -> None:
     )
 
 
-def _check_pyproject() -> None:
+def _update_nbqa_settings() -> None:
+    # cspell:ignore addopts
     if not CONFIG_PATH.precommit.exists():
         return
-    config, _ = load_round_trip_precommit_config()
-    nbqa_repo = find_repo(config, "https://github.com/nbQA-dev/nbQA")
-    if nbqa_repo is None:
+    if not __has_nbqa_precommit_repo():
         return
-    nbqa_config = _load_nbqa_black_config()
-    if nbqa_config != ["--line-length=85"]:
-        error_message = dedent("""
-            Configuration of nbqa-black in pyproject.toml should be as follows:
-
-            [tool.nbqa.addopts]
-            black = [
-                "--line-length=85",
-            ]
-
-            This is to ensure that code blocks render nicely in the sphinx-book-theme.
-            """).strip()
-        raise PrecommitError(error_message)
+    pyproject = load_pyproject()
+    nbqa_table = get_sub_table(pyproject, "tool.nbqa.addopts", create=True)
+    expected = ["--line-length=85"]
+    if nbqa_table.get("black") != expected:
+        nbqa_table["black"] = expected
+        write_pyproject(pyproject)
+        msg = "Added nbQA configuration for black"
+        raise PrecommitError(msg)
 
 
-def _load_nbqa_black_config(content: Optional[str] = None) -> List[str]:
-    # cspell:ignore addopts
-    config = load_pyproject(content)
-    return config.get("tool", {}).get("nbqa", {}).get("addopts", {}).get("black", {})
+def __has_nbqa_precommit_repo() -> bool:
+    config, _ = load_round_trip_precommit_config()
+    nbqa_repo = find_repo(config, "https://github.com/nbQA-dev/nbQA")
+    if nbqa_repo is None:
+        return False
+    return True

src/repoma/utilities/pyproject.py

@@ -1,14 +1,46 @@
-"""Tools for loading and inspecting :code:`pyproject.toml`."""
-from collections import OrderedDict
-from typing import Optional
+"""Tools for loading, inspecting, and updating :code:`pyproject.toml`."""
+import os
+from typing import Any, Iterable, Optional
 
-import toml
+import tomlkit
+from tomlkit.container import Container
+from tomlkit.items import Array, Table
+from tomlkit.toml_document import TOMLDocument
 
 from repoma.utilities import CONFIG_PATH
 
 
-def load_pyproject(content: Optional[str] = None) -> dict:
+def load_pyproject(content: Optional[str] = None) -> TOMLDocument:
+    if not os.path.exists(CONFIG_PATH.pyproject):
+        return TOMLDocument()
     if content is None:
         with open(CONFIG_PATH.pyproject) as stream:
-            return toml.load(stream, _dict=OrderedDict)
-    return toml.loads(content, _dict=OrderedDict)
+            return tomlkit.loads(stream.read())
+    return tomlkit.loads(content)
+
+
+def get_sub_table(config: Container, dotted_header: str, create: bool = False) -> Table:
+    """Get a TOML sub-table through a dotted header key."""
+    current_table: Any = config
+    for header in dotted_header.split("."):
+        if header not in current_table:
+            if create:
+                current_table[header] = tomlkit.table()
+            else:
+                raise KeyError(f"TOML data does not contain {dotted_header!r}")
+        current_table = current_table[header]
+    return current_table
+
+
+def write_pyproject(config: TOMLDocument) -> None:
+    src = tomlkit.dumps(config, sort_keys=True)
+    with open(CONFIG_PATH.pyproject, "w") as stream:
+        stream.write(src)
+
+
+def to_toml_array(items: Iterable[Any]) -> Array:
+    array = tomlkit.array()
+    array.extend(items)
+    if len(array) > 1:
+        array.multiline(True)
+    return array