From 93d7c09dc60fb17c9a216c88ef250d70f3e636f4 Mon Sep 17 00:00:00 2001
From: Jacob Floyd <cognifloyd@gmail.com>
Date: Sat, 19 Mar 2022 07:03:33 -0500
Subject: [PATCH] --write: Introduce TransformMixin for Rules (#2023)

---
 src/ansiblelint/rules/__init__.py | 40 +++++++++++++++++++++++++++++++
 src/ansiblelint/transformer.py    | 27 ++++++++++++++++++---
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/src/ansiblelint/rules/__init__.py b/src/ansiblelint/rules/__init__.py
index 4d6c80848d..0ba739f44f 100644
--- a/src/ansiblelint/rules/__init__.py
+++ b/src/ansiblelint/rules/__init__.py
@@ -12,6 +12,8 @@
 from importlib.abc import Loader
 from typing import Any, Dict, Iterator, List, Optional, Set, Union
 
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
+
 import ansiblelint.skip_utils
 import ansiblelint.utils
 import ansiblelint.yaml_utils
@@ -200,6 +202,44 @@ def matchyaml(self, file: Lintable) -> List[MatchError]:
         return matches
 
 
+class TransformMixin:  # pylint: disable=too-few-public-methods
+    """A mixin for AnsibleLintRule to enable transforming files.
+
+    If ansible-lint is started with the ``--write`` option, then the ``Transformer``
+    will call the ``transform()`` method for every MatchError identified if the rule
+    that identified it subclasses this ``TransformMixin``. Only the rule that identified
+    a MatchError can do transforms to fix that match.
+    """
+
+    def transform(
+        self,
+        match: MatchError,
+        lintable: Lintable,
+        data: Union[CommentedMap, CommentedSeq, str],
+    ) -> None:
+        """Transform ``data`` to try to fix the MatchError identified by this rule.
+
+        The ``match`` was generated by this rule in the ``lintable`` file.
+        When ``transform()`` is called on a rule, the rule should either fix the
+        issue, if possible, or make modifications that make it easier to fix manually.
+
+        For YAML files, ``data`` is an editable YAML dict/array that preserves
+        any comments that were in the original file.
+
+        .. code:: python
+
+            data[0]["tasks"][0]["when"] = False
+
+        For any files that aren't YAML, data is the loaded file's content as a string.
+        To edit non-YAML files, save the updated contents update ``lintable.content``:
+
+        .. code:: python
+
+            new_data = self.do_something_to_fix_the_match(data)
+            lintable.content = new_data
+        """
+
+
 def is_valid_rule(rule: Any) -> bool:
     """Check if given rule is valid or not."""
     return issubclass(rule, AnsibleLintRule) and bool(rule.id) and bool(rule.shortdesc)
diff --git a/src/ansiblelint/transformer.py b/src/ansiblelint/transformer.py
index e4b6391efc..104a2435bd 100644
--- a/src/ansiblelint/transformer.py
+++ b/src/ansiblelint/transformer.py
@@ -1,11 +1,12 @@
 """Transformer implementation."""
 import logging
-from typing import Dict, List, Set, Union
+from typing import Dict, List, Optional, Set, Union
 
 from ruamel.yaml.comments import CommentedMap, CommentedSeq
 
 from ansiblelint.errors import MatchError
 from ansiblelint.file_utils import Lintable
+from ansiblelint.rules import TransformMixin
 from ansiblelint.runner import LintResult
 from ansiblelint.yaml_utils import FormattedYAML
 
@@ -50,7 +51,7 @@ def __init__(self, result: LintResult):
 
     def run(self) -> None:
         """For each file, read it, execute transforms on it, then write it."""
-        for file, _ in self.matches_per_file.items():
+        for file, matches in self.matches_per_file.items():
             # str() convinces mypy that "text/yaml" is a valid Literal.
             # Otherwise, it thinks base_kind is one of playbook, meta, tasks, ...
             file_is_yaml = str(file.base_kind) == "text/yaml"
@@ -62,17 +63,37 @@ def run(self) -> None:
                 data = ""
                 file_is_yaml = False
 
+            ruamel_data: Optional[Union[CommentedMap, CommentedSeq]] = None
             if file_is_yaml:
                 # We need a fresh YAML() instance for each load because ruamel.yaml
                 # stores intermediate state during load which could affect loading
                 # any other files. (Based on suggestion from ruamel.yaml author)
                 yaml = FormattedYAML()
 
-                ruamel_data: Union[CommentedMap, CommentedSeq] = yaml.loads(data)
+                ruamel_data = yaml.loads(data)
                 if not isinstance(ruamel_data, (CommentedMap, CommentedSeq)):
                     # This is an empty vars file or similar which loads as None.
                     # It is not safe to write this file or data-loss is likely.
                     # Only maps and sequences can preserve comments. Skip it.
                     continue
+
+            self._do_transforms(file, ruamel_data or data, matches)
+
+            if file_is_yaml:
+                # noinspection PyUnboundLocalVariable
                 file.content = yaml.dumps(ruamel_data)
+
+            if file.updated:
                 file.write()
+
+    @staticmethod
+    def _do_transforms(
+        file: Lintable,
+        data: Union[CommentedMap, CommentedSeq, str],
+        matches: List[MatchError],
+    ) -> None:
+        """Do Rule-Transforms handling any last-minute MatchError inspections."""
+        for match in sorted(matches):
+            if not isinstance(match.rule, TransformMixin):
+                continue
+            match.rule.transform(match, file, data)