From 58be8ebd77817db3f0881281d401d9c49b1b8f71 Mon Sep 17 00:00:00 2001
From: nbari <nbari@tequila.io>
Date: Thu, 12 Oct 2023 20:17:09 +0200
Subject: [PATCH] preserve ansible # jinja2 header

---
 pyproject.toml              |  2 +-
 src/yamlfix/services.py     |  9 ++++++++-
 tests/unit/test_services.py | 14 ++++++++++++++
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 491d352..a2b88ab 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -69,7 +69,7 @@ source-includes = ["tests/"]
 [tool.pdm.build]
 editable-backend = "path"
 
-[tool.pdm.overrides]
+[tool.pdm.resolution.overrides]
 
 # To be removed once https://github.com/flakeheaven/flakeheaven/issues/132 is solved
 "importlib-metadata" = ">=3.10"
diff --git a/src/yamlfix/services.py b/src/yamlfix/services.py
index e08e52f..7b2ebbe 100644
--- a/src/yamlfix/services.py
+++ b/src/yamlfix/services.py
@@ -149,9 +149,16 @@ def fix_code(source_code: str, config: Optional[YamlfixConfig] = None) -> str:
     else:
         shebang = ""
 
+    if source_code.startswith("#jinja2:") or source_code.startswith("# jinja2:"):
+        eolpos = source_code.find("\n") + 1
+        jinja2 = source_code[:eolpos]
+        source_code = source_code[eolpos:]
+    else:
+        jinja2 = ""
+
     yaml = Yaml(config=config)
     fixer = SourceCodeFixer(yaml=yaml, config=config)
 
     source_code = fixer.fix(source_code=source_code)
 
-    return shebang + source_code
+    return jinja2 + shebang + source_code
diff --git a/tests/unit/test_services.py b/tests/unit/test_services.py
index cbf6bf0..ca23f1b 100644
--- a/tests/unit/test_services.py
+++ b/tests/unit/test_services.py
@@ -74,6 +74,20 @@ def test_fix_files_issues_warning(self, tmp_path: Path) -> None:
 class TestFixCode:
     """Test the fix_code function."""
 
+    def test_fix_code_ignore_jinja2(self) -> None:
+        """Ignores jinja2 line if present at the beginning of the source."""
+        source = dedent(
+            """\
+            # jinja2:lstrip_blocks: true
+            ---
+            program: yamlfix
+            """
+        )
+
+        result = fix_code(source)
+
+        assert result == source
+
     def test_fix_code_ignore_shebang(self) -> None:
         """Ignores shebang lines if present at the beginning of the source."""
         source = dedent(