feat: autofix YAML (recursive madness)

.prettierignore

@@ -5,4 +5,4 @@ _includes/*
 README.md
 
 # Prettier uses double quotes, ruamel.yaml uses single quotes
-tests/test_yaml/2-expected.yaml
+tests/test_yaml/*.yaml

src/nitpick/plugins/yaml.py

@@ -1,32 +1,51 @@
 """YAML files."""
+from collections import OrderedDict
 from itertools import chain
-from typing import Iterator, Optional, Type
-
-from ruamel.yaml import YAML
+from typing import Iterator, List, Optional, Type, Union, cast
 
 from nitpick.constants import PRE_COMMIT_CONFIG_YAML
 from nitpick.formats import BaseFormat, YamlFormat
 from nitpick.plugins import hookimpl
 from nitpick.plugins.base import NitpickPlugin
 from nitpick.plugins.info import FileInfo
 from nitpick.plugins.text import KEY_CONTAINS
-from nitpick.typedefs import YamlObject
+from nitpick.typedefs import JsonDict, YamlObject, YamlValue
 from nitpick.violations import Fuss, SharedViolations, ViolationEnum
 
 
-def change_yaml(document: YAML, dictionary: YamlObject):
+def is_scalar(value: YamlValue) -> bool:
+    """Return True if the value is NOT a dict or a list."""
+    return not isinstance(value, (OrderedDict, list))
+
+
+def traverse_yaml_tree(yaml_obj: YamlObject, change: JsonDict):
     """Traverse a YAML document recursively and change values, keeping its formatting and comments."""
-    del document
-    del dictionary
-    # FIXME:
-    # for key, value in dictionary.items():
-    #     if isinstance(value, (dict, OrderedDict)):
-    #         if key in document:
-    #             change_toml(document[key], value)
-    #         else:
-    #             document.add(key, value)
-    #     else:
-    #         document[key] = value
+    for key, value in change.items():
+        if key not in yaml_obj:
+            # Key doesn't exist: we can insert the whole nested OrderedDict at once, no regrets
+            last_pos = len(yaml_obj.keys()) + 1
+            yaml_obj.insert(last_pos, key, value)
+            continue
+
+        if is_scalar(value):
+            yaml_obj[key] = value
+        elif isinstance(value, OrderedDict):
+            traverse_yaml_tree(yaml_obj[key], value)
+        elif isinstance(value, list):
+            _traverse_yaml_list(yaml_obj, key, value)
+
+
+def _traverse_yaml_list(yaml_obj: YamlObject, key: str, value: List[Union[OrderedDict, str, float]]):
+    for index, element in enumerate(value):
+        if is_scalar(element):
+            try:
+                yaml_obj[key][index] = element
+            except IndexError:
+                yaml_obj[key].append(element)
+        elif isinstance(element, list):
+            pass  # FIXME: a list of lists in YAML? is it possible?
+        else:
+            traverse_yaml_tree(yaml_obj[key][index], cast(OrderedDict, element))
 
 
 class YamlPlugin(NitpickPlugin):
@@ -48,20 +67,19 @@ def enforce_rules(self) -> Iterator[Fuss]:
         if not comparison.has_changes:
             return
 
-        document = yaml_format.document 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, yaml_format.as_object, comparison.diff),
+            self.report(SharedViolations.MISSING_VALUES, yaml_format.as_object, comparison.missing),
         )
-        if self.autofix and self.dirty and document:
-            pass  # FIXME: document.dump(document, self.file_path)
+        if self.autofix and self.dirty:
+            yaml_format.document.dump(yaml_format.as_object, self.file_path)
 
-    def report(self, violation: ViolationEnum, document: Optional[YAML], change: Optional[BaseFormat]):
+    def report(self, violation: ViolationEnum, yaml_object: YamlObject, change: Optional[BaseFormat]):
         """Report a violation while optionally modifying the YAML document."""
         if not change:
             return
-        if document:
-            change_yaml(document, change.as_object)
+        if self.autofix:
+            traverse_yaml_tree(yaml_object, change.as_object)
             self.dirty = True
         yield self.reporter.make_fuss(violation, change.reformatted.strip(), prefix="", fixed=self.autofix)
 

src/nitpick/typedefs.py

@@ -1,4 +1,5 @@
 """Type definitions."""
+from collections import OrderedDict
 from pathlib import Path
 from typing import Any, Dict, Iterable, List, Tuple, Type, Union
 
@@ -10,6 +11,7 @@
 StrOrIterable = Union[str, Iterable[str]]
 Flake8Error = Tuple[int, int, str, Type]
 YamlObject = Union[CommentedSeq, CommentedMap]
+YamlValue = Union[OrderedDict, list, str, float]
 
 # Decorated property not supported · Issue #1362 · python/mypy
 # https://github.com/python/mypy/issues/1362#issuecomment-562141376

tests/helpers.py

@@ -200,6 +200,7 @@ def save_file(self, filename: PathOrStr, file_contents: PathOrStr, lint: bool =
         if lint or path.suffix == ".py":
             self.files_to_lint.append(path)
         clean = dedent(from_path_or_str(file_contents)).strip()
+        path.parent.mkdir(parents=True, exist_ok=True)
         path.write_text(f"{clean}\n")
         return self
 

tests/test_ini.py

@@ -40,7 +40,7 @@ def test_default_style_is_applied(project_default):
         strict_optional = True
         warn_no_return = True
         warn_redundant_casts = True
-        warn_unused_ignores = True
+        warn_unused_ignores = False
     """
     expected_editor_config = """
         root = True

tests/test_yaml.py

@@ -7,8 +7,8 @@
 def test_suggest_initial_contents(tmp_path, datadir):
     """Suggest contents when YAML files do not exist."""
     filename = ".github/workflows/python.yaml"
-    expected_yaml = (datadir / "2-expected.yaml").read_text()
-    ProjectMock(tmp_path).style(datadir / "2-desired.toml").api_check_then_fix(
+    expected_yaml = (datadir / "new-expected.yaml").read_text()
+    ProjectMock(tmp_path).style(datadir / "new-desired.toml").api_check_then_fix(
         Fuss(
             True,
             filename,
@@ -21,8 +21,10 @@ def test_suggest_initial_contents(tmp_path, datadir):
 
 def test_missing_different_values(tmp_path, datadir):
     """Test different and missing values on any YAML."""
-    filename = "1-actual.yaml"
-    ProjectMock(tmp_path).save_file(filename, datadir / filename).style(datadir / "1-desired.toml").api_check_then_fix(
+    filename = "me/deep/rooted.yaml"
+    ProjectMock(tmp_path).save_file(filename, datadir / "existing-actual.yaml").style(
+        datadir / "existing-desired.toml"
+    ).api_check_then_fix(
         Fuss(
             True,
             filename,
@@ -32,9 +34,9 @@ def test_missing_different_values(tmp_path, datadir):
             python:
               install:
                 - extra_requirements:
-                    - item1
-                    - item2
-                    - item3
+                    - some
+                    - nice
+                    - package
               version: '3.9'
             """,
         ),
@@ -46,15 +48,17 @@ def test_missing_different_values(tmp_path, datadir):
             """
             root_key:
               a_dict:
-                - a: string value
-                - b: 2
                 - c: '3.1'
+                - b: 2
+                - a: string value
               a_nested:
                 int: 10
                 list:
                   - 0
-                  - 1
                   - 2
+                  - 1
             """,
         ),
-    )  # FIXME:  .assert_file_contents(filename, datadir / "1-expected.yaml")
+    ).assert_file_contents(
+        filename, datadir / "existing-expected.yaml"
+    )

tests/test_yaml/1-actual.yaml → tests/test_yaml/existing-actual.yaml

@@ -1,5 +1,6 @@
+# Root comment
 python:
-  version: 3.6
+  version: 3.6  # Python 3.6 EOL this month!
   install:
     - method: pip
       path: .

tests/test_yaml/existing-desired.toml

@@ -0,0 +1,18 @@
+["me/deep/rooted.yaml".python]
+version = "3.9"
+
+[["me/deep/rooted.yaml".python.install]]
+extra_requirements = ["some", "nice", "package"]
+
+[["me/deep/rooted.yaml".root_key.a_dict]]
+c = "3.1"
+
+[["me/deep/rooted.yaml".root_key.a_dict]]
+b = 2
+
+[["me/deep/rooted.yaml".root_key.a_dict]]
+a = "string value"
+
+["me/deep/rooted.yaml".root_key.a_nested]
+list = [0, 2, 1]
+int = 10

tests/test_yaml/existing-expected.yaml

@@ -0,0 +1,21 @@
+# Root comment
+python:
+  version: '3.9' # Python 3.6 EOL this month!
+  install:
+    - method: pip
+      path: .
+      extra_requirements:
+        - some
+        - nice
+        - package
+root_key:
+  a_dict:
+    - c: '3.1'
+    - b: 2
+    - a: string value
+  a_nested:
+    int: 10
+    list:
+      - 0
+      - 2
+      - 1