feat: TomlTable helper class

src/nitpick/blender.py

@@ -10,6 +10,7 @@
 import json
 import re
 import shlex
+from dataclasses import dataclass
 from functools import partial
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast
@@ -23,7 +24,7 @@
 from flatten_dict import flatten, unflatten
 from ruamel.yaml import YAML, RoundTripRepresenter, StringIO
 from sortedcontainers import SortedDict
-from tomlkit import items
+from tomlkit import TOMLDocument, items
 
 from nitpick.typedefs import ElementData, JsonDict, ListOrCommentedSeq, PathOrStr, YamlObject, YamlValue
 
@@ -553,6 +554,71 @@ def load(self) -> bool:
         return True
 
 
+@dataclass
+class TomlTable:
+    """A helper to write data on a TOML table using tomlkit to preserve existing content."""
+
+    path: Path
+    name: str
+    _doc: TOMLDocument = None
+    _table: items.Table = None
+    _created: bool = False
+
+    def _parse_file(self):
+        if self._doc is None:
+            if self.path.exists():
+                self._doc = tomlkit.loads(self.path.read_text(encoding="UTF-8"))
+            else:
+                self._doc = tomlkit.document()
+
+    def _get_table(self, *, create=False) -> items.Table | None:
+        if self._table is not None:
+            return self._table
+
+        self._parse_file()
+        current = self._doc
+        for part in self.name.split(SEPARATOR_DOT):
+            if part not in current:
+                if create:
+                    self._table = tomlkit.table()
+                    self._created = True
+                return self._table
+            current = current[part]
+
+        self._table = current
+        return self._table
+
+    @property
+    def exists(self) -> bool:
+        """Return True if the table exists."""
+        return self._get_table() is not None
+
+    def update_list(self, key: str, *args: str, comment: str = "") -> None:
+        """Update a list on the table with an optional comment."""
+        self._get_table(create=True)
+
+        normalized_array = tomlkit.array([tomlkit.string(item) for item in args])
+        if self._table.get(key):
+            # TODO(AA): add comment or update existing one
+            self._table[key] = normalized_array
+            return
+
+        if comment:
+            self._table.add(tomlkit.comment("\n# ".join(comment.splitlines())))
+        self._table.add(key, normalized_array)
+
+    def write_file(self):
+        """Write the TOML file."""
+        if self._created:
+            self._doc.add(items.SingleKey(self.name, items.KeyType.Bare), self._table)
+        self.path.write_text(tomlkit.dumps(self._doc), encoding="UTF-8")
+
+    @property
+    def as_toml(self) -> str:
+        """Return the table as a TOML string."""
+        return tomlkit.dumps(self._table).rstrip()
+
+
 def traverse_toml_tree(document: tomlkit.TOMLDocument, dictionary):
     """Traverse a TOML document recursively and change values, keeping its formatting and comments."""
     for key, value in dictionary.items():

src/nitpick/cli.py

@@ -19,16 +19,16 @@
 from pathlib import Path
 
 import click
-import tomlkit
 from click.exceptions import Exit
 from loguru import logger
 
-from nitpick.blender import TomlDoc
-from nitpick.constants import PROJECT_NAME, TOOL_KEY, TOOL_NITPICK_KEY
+from nitpick.blender import TomlTable
+from nitpick.constants import DOT_NITPICK_TOML, PROJECT_NAME, READ_THE_DOCS_URL, TOOL_NITPICK_KEY
 from nitpick.core import Nitpick
 from nitpick.enums import OptionEnum
 from nitpick.exceptions import QuitComplainingError
 from nitpick.generic import relative_to_current_dir
+from nitpick.style import StyleManager
 from nitpick.violations import Reporter
 
 VERBOSE_OPTION = click.option(
@@ -156,19 +156,25 @@ def ls(context, files):  # pylint: disable=invalid-name
 def init(context, force, style_urls):
     """Create a [tool.nitpick] section in the configuration file."""
     nit = get_nitpick(context)
+    # TODO(AA): test --force flag
+    # TODO(AA): replace both these lines with path = nit.project.which_config_file()
     config = nit.project.read_configuration()
+    path = config.file or nit.project.root / DOT_NITPICK_TOML
 
-    if config.file:
-        tool_nitpick_toml = TomlDoc(path=config.file).as_object.get(TOOL_KEY, {}).get(PROJECT_NAME, {})
-        if tool_nitpick_toml and not force:
-            click.secho(f"{PROJECT_NAME} is already configured in {config.file.name!r}.", fg="yellow")
-            click.echo(f"[{TOOL_NITPICK_KEY}]")
-            click.echo(tomlkit.dumps(tool_nitpick_toml))
-            raise Exit(1)
-
-    nit.project.create_configuration(config, *style_urls)
-    click.secho(f"A [{TOOL_NITPICK_KEY}] section was created in the config file: {config.file.name}", fg="green")
+    table = TomlTable(path, TOOL_NITPICK_KEY)
+    if table.exists and not force:
+        click.secho(f"The config file {path.name!r} already has a [{TOOL_NITPICK_KEY}] section.", fg="yellow")
+        click.echo(table.as_toml)
+        raise Exit(1)
 
-    tool_nitpick_toml = TomlDoc(path=config.file).as_object.get(TOOL_KEY, {}).get(PROJECT_NAME, {})
-    click.echo(f"[{TOOL_NITPICK_KEY}]")
-    click.echo(tomlkit.dumps(tool_nitpick_toml))
+    if not style_urls:
+        style_urls = (str(StyleManager.get_default_style_url()),)
+    table.update_list(
+        "style",
+        *style_urls,
+        comment=f"Generated by the 'nitpick init' command\nMore info at {READ_THE_DOCS_URL}configuration.html",
+    )
+    table.write_file()
+    verb = "updated" if force else "created"
+    click.secho(f"The [{TOOL_NITPICK_KEY}] section was {verb} in the config file {path.name!r}", fg="green")
+    click.echo(table.as_toml)

src/nitpick/project.py

@@ -6,25 +6,20 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Iterable, Iterator
 
-import tomlkit
 from autorepr import autorepr
 from loguru import logger
 from marshmallow_polyfield import PolyField
 from more_itertools.more import always_iterable
 from pluggy import PluginManager
-from tomlkit.exceptions import NonExistentKey
-from tomlkit.items import KeyType, SingleKey
 
 from nitpick import fields, plugins
 from nitpick.blender import TomlDoc, search_json
 from nitpick.constants import (
     CONFIG_FILES,
-    DOT_NITPICK_TOML,
     MANAGE_PY,
     NITPICK_MINIMUM_VERSION_JMEX,
     PROJECT_NAME,
     PYPROJECT_TOML,
-    READ_THE_DOCS_URL,
     ROOT_FILES,
     ROOT_PYTHON_FILES,
     TOOL_NITPICK_JMEX,
@@ -36,8 +31,6 @@
 from nitpick.violations import Fuss, ProjectViolations, Reporter, StyleViolations
 
 if TYPE_CHECKING:
-    from tomlkit.toml_document import TOMLDocument
-
     from nitpick.typedefs import JsonDict, PathOrStr
 
 
@@ -205,38 +198,3 @@ def merge_styles(self, offline: bool) -> Iterator[Fuss]:
 
         self.nitpick_section = self.style_dict.get("nitpick", {})
         self.nitpick_files_section = self.nitpick_section.get("files", {})
-
-    def create_configuration(self, config: Configuration, *style_urls: str) -> str:
-        """Create a configuration file."""
-        from nitpick.style import StyleManager  # pylint: disable=import-outside-toplevel
-
-        if config.file:
-            doc: TOMLDocument = tomlkit.parse(config.file.read_text())
-        else:
-            doc = tomlkit.document()
-            config.file = self.root / DOT_NITPICK_TOML
-
-        if not style_urls:
-            style_urls = (str(StyleManager.get_default_style_url()),)
-
-        try:
-            tool_nitpick = doc["tool"]["nitpick"]
-            section_exists = True
-        except NonExistentKey:
-            section_exists = False
-            tool_nitpick = tomlkit.table()
-            tool_nitpick.add(tomlkit.comment("Generated by the 'nitpick init' command"))
-            tool_nitpick.add(tomlkit.comment(f"More info at {READ_THE_DOCS_URL}configuration.html"))
-
-        styles = tomlkit.array([tomlkit.string(url) for url in style_urls])
-        if tool_nitpick.get("style"):
-            tool_nitpick["style"] = styles
-        else:
-            tool_nitpick.add("style", styles)
-
-        if not section_exists:
-            doc.add(SingleKey(TOOL_NITPICK_KEY, KeyType.Bare), tool_nitpick)
-
-        rv = tomlkit.dumps(doc, sort_keys=True)
-        config.file.write_text(rv)
-        return rv

tests/test_cli.py

@@ -45,15 +45,21 @@ def test_config_file_already_has_tool_nitpick_section(tmp_path, config_file):
         style = ['/this/should/not/be/validated-yet.toml']
         """,
     )
-    project.cli_init(f"The config file {config_file} already has a [{TOOL_NITPICK_KEY}] section.", exit_code=1)
+    project.cli_init(
+        [
+            f"The config file {config_file!r} already has a [{TOOL_NITPICK_KEY}] section.",
+            "style = ['/this/should/not/be/validated-yet.toml']",
+        ],
+        exit_code=1,
+    )
 
 
 def test_create_basic_dot_nitpick_toml(tmp_path):
     """If no config file is found, create a basic .nitpick.toml."""
     project = ProjectMock(tmp_path, pyproject_toml=False, setup_py=True)
     url = StyleManager.get_default_style_url()
     project.cli_init(
-        f"A [{TOOL_NITPICK_KEY}] section was created in the config file: {DOT_NITPICK_TOML}"
+        f"The [{TOOL_NITPICK_KEY}] section was created in the config file {DOT_NITPICK_TOML!r}\nstyle = ['{url}']"
     ).assert_file_contents(
         DOT_NITPICK_TOML,
         f"""
@@ -71,7 +77,10 @@ def test_init_empty_pyproject_toml(tmp_path):
     project = ProjectMock(tmp_path, pyproject_toml=False, setup_py=True)
     url = StyleManager.get_default_style_url()
     project.pyproject_toml("").cli_init(
-        f"A [{TOOL_NITPICK_KEY}] section was created in the config file: {PYPROJECT_TOML}"
+        [
+            f"The [{TOOL_NITPICK_KEY}] section was created in the config file {PYPROJECT_TOML!r}",
+            "style = ['py://nitpick/resources/presets/nitpick']",
+        ]
     ).assert_file_contents(
         PYPROJECT_TOML,
         f"""
@@ -86,13 +95,16 @@ def test_init_empty_pyproject_toml(tmp_path):
 
 
 @pytest.mark.parametrize(
-    "styles",
+    ("styles", "expected_styles"),
     [
-        (),  # no arguments, default style
-        ("https://github.com/andreoliwa/nitpick/blob/develop/initial.toml", "./local.toml"),
+        ((), "style = ['py://nitpick/resources/presets/nitpick']"),  # no arguments, default style
+        (
+            ("https://github.com/andreoliwa/nitpick/blob/develop/initial.toml", "./local.toml"),
+            "style = ['https://github.com/andreoliwa/nitpick/blob/develop/initial.toml', './local.toml']",
+        ),
     ],
 )
-def test_add_tool_nitpick_section_to_pyproject_toml(tmp_path, styles):
+def test_add_tool_nitpick_section_to_pyproject_toml(tmp_path, styles, expected_styles):
     """Add a [tool.nitpick] section to pyproject.toml."""
     project = ProjectMock(tmp_path).pyproject_toml(
         """
@@ -103,7 +115,8 @@ def test_add_tool_nitpick_section_to_pyproject_toml(tmp_path, styles):
     expected = styles or [StyleManager.get_default_style_url()]
 
     project.cli_init(
-        f"A [{TOOL_NITPICK_KEY}] section was created in the config file: {PYPROJECT_TOML}", *styles
+        f"The [{TOOL_NITPICK_KEY}] section was created in the config file {PYPROJECT_TOML!r}\n{expected_styles}",
+        *styles,
     ).assert_file_contents(
         PYPROJECT_TOML,
         f"""