Skip to content

Commit

Permalink
feat: don't erase comment-only yaml documents
Browse files Browse the repository at this point in the history
  • Loading branch information
marcules committed Dec 26, 2022
1 parent 867880d commit 3ab4d87
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 1 deletion.
34 changes: 34 additions & 0 deletions src/yamlfix/adapters.py
Expand Up @@ -2,6 +2,7 @@

import logging
import re
import uuid
from io import StringIO
from typing import Any, Callable, List, Optional, Tuple

Expand Down Expand Up @@ -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.
Expand All @@ -344,13 +346,15 @@ 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,
self._ruamel_yaml_fixer,
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,
Expand Down Expand Up @@ -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)
34 changes: 33 additions & 1 deletion tests/unit/test_services.py
Expand Up @@ -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(
Expand All @@ -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(
Expand Down

0 comments on commit 3ab4d87

Please sign in to comment.