From dbd9ac7d0daf479447675a696392424786fe6302 Mon Sep 17 00:00:00 2001
From: Dany Gielow <dany.gielow@gmail.com>
Date: Mon, 29 May 2023 17:01:48 +0200
Subject: [PATCH] fix: changed TOML lists are replaced (same behaviour as in
 YAML files) (#577)

---
 src/nitpick/plugins/toml.py | 32 +++++++++++++++++++++++---------
 tests/test_toml.py          | 17 +++++++++++++++++
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/src/nitpick/plugins/toml.py b/src/nitpick/plugins/toml.py
index 6db9fd11..4c88d759 100644
--- a/src/nitpick/plugins/toml.py
+++ b/src/nitpick/plugins/toml.py
@@ -2,11 +2,11 @@
 from __future__ import annotations
 
 from itertools import chain
-from typing import TYPE_CHECKING, Iterator
+from typing import TYPE_CHECKING, Iterator, cast
 
 from tomlkit import dumps, parse
 
-from nitpick.blender import BaseDoc, Comparison, TomlDoc, traverse_toml_tree
+from nitpick.blender import Comparison, TomlDoc, traverse_toml_tree
 from nitpick.plugins import hookimpl
 from nitpick.plugins.base import NitpickPlugin
 from nitpick.violations import Fuss, SharedViolations, ViolationEnum
@@ -42,20 +42,34 @@ def enforce_rules(self) -> Iterator[Fuss]:
 
         document = parse(toml_doc.as_string) if self.autofix else None
         yield from chain(
-            self.report(SharedViolations.DIFFERENT_VALUES, document, comparison.diff),
-            self.report(SharedViolations.MISSING_VALUES, document, comparison.missing),
+            self.report(SharedViolations.DIFFERENT_VALUES, document, cast(TomlDoc, comparison.diff)),
+            self.report(
+                SharedViolations.MISSING_VALUES,
+                document,
+                cast(TomlDoc, comparison.missing),
+                cast(TomlDoc, comparison.replace),
+            ),
         )
         if self.autofix and self.dirty:
             self.file_path.write_text(dumps(document))
 
-    def report(self, violation: ViolationEnum, document: TOMLDocument | None, change: BaseDoc | None):
+    def report(
+        self,
+        violation: ViolationEnum,
+        document: TOMLDocument | None,
+        change: TomlDoc | None,
+        replacement: TomlDoc | None = None,
+    ):
         """Report a violation while optionally modifying the TOML document."""
-        if not change:
+        if not (change or replacement):
             return
-        if document:
-            traverse_toml_tree(document, change.as_object)
+        if self.autofix:
+            real_change = cast(TomlDoc, replacement or change)
+            traverse_toml_tree(document, real_change.as_object)
             self.dirty = True
-        yield self.reporter.make_fuss(violation, change.reformatted.strip(), prefix="", fixed=self.autofix)
+
+        to_display = cast(TomlDoc, change or replacement)
+        yield self.reporter.make_fuss(violation, to_display.reformatted.strip(), prefix="", fixed=self.autofix)
 
     @property
     def initial_contents(self) -> str:
diff --git a/tests/test_toml.py b/tests/test_toml.py
index 432d7594..0b339727 100644
--- a/tests/test_toml.py
+++ b/tests/test_toml.py
@@ -29,12 +29,14 @@ def test_suggest_initial_contents(tmp_path):
         [section]
         key = "value"
         number = 10
+        list = [ "a", "b", "c",]
     """
     ProjectMock(tmp_path).style(
         f"""
         ["{filename}".section]
         key = "value"
         number = 10
+        list = ["a", "b", "c"]
         """
     ).api_check_then_fix(
         Fuss(
@@ -58,6 +60,9 @@ def test_missing_different_values_pyproject_toml(tmp_path):
 
         ["pyproject.toml".tool]
         missing = "value"
+
+        ["pyproject.toml".config]
+        list = ["a", "b", "c"]
         """
     ).pyproject_toml(
         """
@@ -65,6 +70,9 @@ def test_missing_different_values_pyproject_toml(tmp_path):
         x = 1  # comment for x
         yada = "before"  # comment for yada yada
         abc = "123" # comment for abc
+
+        [config]
+        list = ["a", "b"] # comment for list
         """
     ).api_check_then_fix(
         Fuss(
@@ -85,6 +93,9 @@ def test_missing_different_values_pyproject_toml(tmp_path):
             """
             [tool]
             missing = "value"
+
+            [config]
+            list = [ "c",]
             """,
         ),
     ).assert_file_contents(
@@ -95,6 +106,9 @@ def test_missing_different_values_pyproject_toml(tmp_path):
         yada = "after"  # comment for yada yada
         abc = "123" # comment for abc
 
+        [config]
+        list = ["a", "b", "c"] # comment for list
+
         [tool]
         missing = "value"
         """,
@@ -116,6 +130,7 @@ def test_missing_different_values_any_toml(tmp_path):
         ["{filename}".section]
         key = "new value"
         number = 5
+        list = ["a", "b", "c"]
         """
     ).api_check_then_fix(
         Fuss(
@@ -136,6 +151,7 @@ def test_missing_different_values_any_toml(tmp_path):
             """
             [section]
             number = 5
+            list = [ "a", "b", "c",]
             """,
         ),
     ).assert_file_contents(
@@ -145,5 +161,6 @@ def test_missing_different_values_any_toml(tmp_path):
         # Line comment
         key = "new value"
         number = 5
+        list = ["a", "b", "c"]
         """,
     )