From 3ab4d870aaec0f79abe9f56df02fd9affb5e71ef Mon Sep 17 00:00:00 2001
From: Marc Harnos <mharnos@gmail.com>
Date: Tue, 27 Dec 2022 00:48:03 +0100
Subject: [PATCH] feat: don't erase comment-only yaml documents

---
 src/yamlfix/adapters.py     | 34 ++++++++++++++++++++++++++++++++++
 tests/unit/test_services.py | 34 +++++++++++++++++++++++++++++++++-
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/src/yamlfix/adapters.py b/src/yamlfix/adapters.py
index fbb9dcc..40b2d06 100644
--- a/src/yamlfix/adapters.py
+++ b/src/yamlfix/adapters.py
@@ -2,6 +2,7 @@
 
 import logging
 import re
+import uuid
 from io import StringIO
 from typing import Any, Callable, List, Optional, Tuple
 
@@ -331,6 +332,7 @@ def __init__(self, yaml: Yaml, config: Optional[YamlfixConfig]) -> None:
         """
         self.yaml = yaml.yaml
         self.config = config or YamlfixConfig()
+        self.document_fix_id = str(uuid.uuid4())
 
     def fix(self, source_code: str) -> str:
         """Run all yaml source code fixers.
@@ -344,6 +346,7 @@ def fix(self, source_code: str) -> str:
         log.debug("Running source code fixers...")
 
         fixers = [
+            self._fix_comment_only_files,
             self._fix_truthy_strings,
             self._fix_comments,
             self._fix_jinja_variables,
@@ -351,6 +354,7 @@ def fix(self, source_code: str) -> str:
             self._restore_truthy_strings,
             self._restore_double_exclamations,
             self._restore_jinja_variables,
+            self._restore_comment_only_files,
             self._fix_top_level_lists,
             self._fix_flow_style_lists,
             self._add_newline_at_end_of_file,
@@ -689,3 +693,33 @@ def _restore_jinja_variables(source_code: str) -> str:
             fixed_source_lines.append(line)
 
         return "\n".join(fixed_source_lines)
+
+    def _fix_comment_only_files(self, source_code: str) -> str:
+        """Add a mapping key with an id to the start of the document\
+            to preserve comments."""
+        fixed_source_lines = []
+        yamlfix_document_id_line = f"yamlfix_document_fix_id:  # {self.document_fix_id}"
+
+        # Add the document id line after each document start
+        has_start_indicator = False
+        for line in source_code.splitlines():
+            fixed_source_lines.append(line)
+            if "---" in line:
+                has_start_indicator = True
+                fixed_source_lines.append(yamlfix_document_id_line)
+
+        # if the document has no start indicator, the document id as the first line
+        if not has_start_indicator:
+            fixed_source_lines.insert(0, yamlfix_document_id_line)
+
+        return "\n".join(fixed_source_lines)
+
+    def _restore_comment_only_files(self, source_code: str) -> str:
+        """Remove the document start id from the document again."""
+        fixed_source_lines = []
+
+        for line in source_code.splitlines():
+            if self.document_fix_id not in line:
+                fixed_source_lines.append(line)
+
+        return "\n".join(fixed_source_lines)
diff --git a/tests/unit/test_services.py b/tests/unit/test_services.py
index da606a2..bc2c022 100644
--- a/tests/unit/test_services.py
+++ b/tests/unit/test_services.py
@@ -180,6 +180,21 @@ def test_fix_code_preserves_comments(self) -> None:
 
         assert result == source
 
+    def test_fix_code_preserves_comments_without_start_indication(self) -> None:
+        """Don't delete comments without yaml explicit start indictor."""
+        source = dedent(
+            """\
+            # Keep comments!
+            program: yamlfix
+            """
+        )
+        config = YamlfixConfig()
+        config.explicit_start = False
+
+        result = fix_code(source, config)
+
+        assert result == source
+
     def test_fix_code_preserves_comment_only_file(self) -> None:
         """Don't delete comments even if the file is only comments."""
         source = dedent(
@@ -192,7 +207,24 @@ def test_fix_code_preserves_comment_only_file(self) -> None:
         result = fix_code(source)
 
         assert result == source
-        
+
+    def test_fix_code_preserves_comment_only_files_without_start_indication(
+        self,
+    ) -> None:
+        """Don't delete comments even if the file is only comments, without\
+            start indication."""
+        source = dedent(
+            """\
+            # Keep comments!
+            """
+        )
+        config = YamlfixConfig()
+        config.explicit_start = False
+
+        result = fix_code(source, config)
+
+        assert result == source
+
     def test_fix_code_respects_parent_lists_with_comments(self) -> None:
         """Do not indent lists at the first level even if there is a comment."""
         source = dedent(