feat: correct truthy strings (#3)

docs/index.md

@@ -54,6 +54,16 @@ book_library:
     fix_files(['file.py'])
     ```
 
+# Features
+
+yamlfix will do the following changes in your code:
+
+* Add the header `---` to your file.
+* [Correct truthy
+    strings](https://yamllint.readthedocs.io/en/stable/rules.html#module-yamllint.rules.truthy):
+    'True' -> true, 'no' -> 'false'
+* Remove unnecessary apostrophes: `title: 'Why we sleep'` -> `title: Why we sleep`.
+
 # References
 
 As most open sourced programs, `yamlfix` is standing on the shoulders of

src/yamlfix/services.py

@@ -57,14 +57,22 @@ def fix_code(source_code: str) -> str:
     It corrects these errors:
 
         * Add --- at the beginning of the file.
+        * Correct truthy strings: 'True' -> true, 'no' -> 'false'
+        * Remove unnecessary apostrophes: `title: 'Why we sleep'` ->
+            `title: Why we sleep`.
 
     Args:
         source_code: Source code to be corrected.
 
     Returns:
         Corrected source code.
     """
-    fixers = [_ruamel_yaml_fixer, _fix_top_level_lists]
+    fixers = [
+        _fix_truthy_strings,
+        _ruamel_yaml_fixer,
+        _restore_truthy_strings,
+        _fix_top_level_lists,
+    ]
     for fixer in fixers:
         source_code = fixer(source_code)
 
@@ -140,7 +148,7 @@ def _fix_top_level_lists(source_code: str) -> str:
 
             # Extract the indentation level
             serialized_line = re.match(r"(?P<indent>\s*)- +(?P<content>.*)", line)
-            if serialized_line is None:
+            if serialized_line is None:  # pragma: no cover
                 raise ValueError(f"Error extracting the indentation of line: {line}")
             indent = serialized_line.groupdict()["indent"]
 
@@ -152,3 +160,81 @@ def _fix_top_level_lists(source_code: str) -> str:
             return source_code
 
     return "\n".join(fixed_source_lines)
+
+
+def _fix_truthy_strings(source_code: str) -> str:
+    """Convert common strings that refer to booleans.
+
+    All caps variations of true, yes and on are transformed to true, while false, no and
+    off are transformed to false.
+
+    Ruyaml understands these strings and converts them to the lower version of the word
+    instead of converting them to true and false.
+
+    [More
+    info](https://yamllint.readthedocs.io/en/stable/rules.html#module-yamllint.rules.truthy)
+
+    Args:
+        source_code: Source code to be corrected.
+
+    Returns:
+        Corrected source code.
+    """
+    source_lines = source_code.splitlines()
+    fixed_source_lines: List[str] = []
+
+    for line in source_lines:
+        line_contains_true = re.match(
+            r"(?P<pre_boolean_text>.*(:|-) )(true|yes|on)$", line, re.IGNORECASE
+        )
+        line_contains_false = re.match(
+            r"(?P<pre_boolean_text>.*(:|-) )(false|no|off)$", line, re.IGNORECASE
+        )
+
+        if line_contains_true:
+            fixed_source_lines.append(
+                f"{line_contains_true.groupdict()['pre_boolean_text']}true"
+            )
+        elif line_contains_false:
+            fixed_source_lines.append(
+                f"{line_contains_false.groupdict()['pre_boolean_text']}false"
+            )
+        else:
+            fixed_source_lines.append(line)
+
+    return "\n".join(fixed_source_lines)
+
+
+def _restore_truthy_strings(source_code: str) -> str:
+    """Restore truthy strings to strings.
+
+    The Ruyaml parser removes the apostrophes of all the caps variations of the strings
+    'yes', 'on', no and 'off' as it interprets them as booleans.
+
+    As this function is run after _fix_truthy_strings, those strings are meant to be
+    strings. So we're turning them back from booleans to strings.
+
+    Args:
+        source_code: Source code to be corrected.
+
+    Returns:
+        Corrected source code.
+    """
+    source_lines = source_code.splitlines()
+    fixed_source_lines: List[str] = []
+
+    for line in source_lines:
+        line_contains_valid_truthy_string = re.match(
+            r"(?P<pre_boolean_text>.*(:|-) )(?P<boolean_text>yes|on|no|off)$",
+            line,
+            re.IGNORECASE,
+        )
+        if line_contains_valid_truthy_string:
+            fixed_source_lines.append(
+                f"{line_contains_valid_truthy_string.groupdict()['pre_boolean_text']}"
+                f"'{line_contains_valid_truthy_string.groupdict()['boolean_text']}'"
+            )
+        else:
+            fixed_source_lines.append(line)
+
+    return "\n".join(fixed_source_lines)

tests/unit/test_services.py

@@ -2,8 +2,34 @@
 
 from textwrap import dedent
 
+import pytest
+
 from yamlfix.services import fix_code
 
+true_strings = [
+    "TRUE",
+    "True",
+    "true",
+    "YES",
+    "Yes",
+    "yes",
+    "ON",
+    "On",
+    "on",
+]
+
+false_strings = [
+    "FALSE",
+    "False",
+    "false",
+    "NO",
+    "No",
+    "no",
+    "OFF",
+    "Off",
+    "off",
+]
+
 
 def test_fix_code_adds_header() -> None:
     """Adds the --- at the beginning of the source."""
@@ -80,3 +106,94 @@ def test_fix_code_preserves_comments() -> None:
     result = fix_code(source)
 
     assert result == source
+
+
+def test_fix_code_removes_extra_apostrophes() -> None:
+    """Remove not needed apostrophes."""
+    source = dedent(
+        """\
+        ---
+        title: 'Why we sleep'"""
+    )
+    fixed_source = dedent(
+        """\
+        ---
+        title: Why we sleep"""
+    )
+
+    result = fix_code(source)
+
+    assert result == fixed_source
+
+
+@pytest.mark.parametrize("true_string", true_strings)
+def test_fix_code_converts_non_valid_true_booleans(true_string: str) -> None:
+    """Convert common strings that refer to true, but aren't the string `true`.
+
+    [More
+    info](https://yamllint.readthedocs.io/en/stable/rules.html#module-yamllint.rules.truthy)
+    """
+    source = dedent(
+        f"""\
+        ---
+        True dictionary: {true_string}
+        True list:
+          - {true_string}"""
+    )
+    fixed_source = dedent(
+        """\
+        ---
+        True dictionary: true
+        True list:
+          - true"""
+    )
+
+    result = fix_code(source)
+
+    assert result == fixed_source
+
+
+@pytest.mark.parametrize("false_string", false_strings)
+def test_fix_code_converts_non_valid_false_booleans(false_string: str) -> None:
+    """Convert common strings that refer to false, but aren't the string `false`.
+
+    [More
+    info](https://yamllint.readthedocs.io/en/stable/rules.html#module-yamllint.rules.truthy)
+    """
+    source = dedent(
+        f"""\
+        ---
+        False dictionary: {false_string}
+        False list:
+          - {false_string}"""
+    )
+    fixed_source = dedent(
+        """\
+        ---
+        False dictionary: false
+        False list:
+          - false"""
+    )
+
+    result = fix_code(source)
+
+    assert result == fixed_source
+
+
+@pytest.mark.parametrize("truthy_string", true_strings + false_strings)
+def test_fix_code_respects_apostrophes_for_truthy_substitutions(
+    truthy_string: str,
+) -> None:
+    """Keep apostrophes for strings like `yes` or `true`.
+
+    So they are not converted to booleans.
+    """
+    source = dedent(
+        f"""\
+        ---
+        title: '{truthy_string}'"""
+    )
+
+    result = fix_code(source)
+
+    assert result == source