fix/#600/jinja multiline string (#625)

* fix #600: improve detection of python multiline strings * chore: update changelog * fix: also fix case using tuple concat strings
tconbeer · Jul 26, 2024 · adee5d7 · adee5d7
1 parent c1ac50b
commit adee5d7
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+### Bug Fixes
+
+- Fixes a bug where complex python statements inside of jinja tags could cause unstable formatting ([#600](https://github.com/tconbeer/sqlfmt/issues/600) - thank you [@nenkie76](https://github.com/nenkie76)!)
+
 ## [0.23.1] - 2024-07-26
 
 ### Bug Fixes

diff --git a/src/sqlfmt/jinjafmt.py b/src/sqlfmt/jinjafmt.py
@@ -271,9 +271,8 @@ def _multiline_str(self) -> str:
             extra_indent = " " * 4
 
         for i, code_line in enumerate(code_lines, start=1 if self.verb else 0):
-            lines.append(
-                f"{indent}{'' if i in no_indent_lines else extra_indent}{code_line}"
-            )
+            line_indent = "" if i in no_indent_lines else f"{indent}{extra_indent}"
+            lines.append(f"{line_indent}{code_line}")
 
         if self.verb:
             lines[-1] = f"{indent}{lines[-1].lstrip()} {self.closing_marker}"
@@ -286,19 +285,35 @@ def _basic_str(self) -> str:
         return f"{self.opening_marker} {self.verb}{self.code} {self.closing_marker}"
 
     def _find_multiline_python_str_lines(self) -> MutableSet[int]:
+        """
+        Return a set line numbers that correspond with the lines
+        of a triple-quoted multiline string (except for the first line
+        of each string). These are lines that should never have
+        their indentation adjusted
+        """
         try:
-            tree = ast.parse(self.code, mode="eval")
+            tree = ast.parse(self.code, mode="exec")
         except SyntaxError:
             # this jinja isn't quite python, so give up here.
             return set()
 
         line_indicies: MutableSet[int] = set()
+        raw_lines = self.code.splitlines()
         for node in ast.walk(tree):
             if (
                 isinstance(node, ast.Constant)
                 and isinstance(node.value, str)
-                and "\n" in node.value
                 and node.end_lineno is not None
+                and node.end_lineno > node.lineno
+                and "\n" in node.value
+                and (
+                    '"""' in raw_lines[node.lineno - 1]
+                    or "'''" in raw_lines[node.lineno - 1]
+                )
+                and (
+                    '"""' in raw_lines[node.end_lineno - 1]
+                    or "'''" in raw_lines[node.end_lineno - 1]
+                )
             ):
                 line_indicies |= set(range(node.lineno, node.end_lineno))
 

diff --git a/tests/data/preformatted/303_jinjafmt_more_mutliline_str.sql b/tests/data/preformatted/303_jinjafmt_more_mutliline_str.sql
@@ -0,0 +1,18 @@
+{% macro test(key) %}
+    {% set columns_by_source = {
+        "a": """
+        column_a1,
+        column_a2,
+        column_a3,
+        """,
+        "b": """
+        column_b1,
+        column_b2,
+        """,
+        "c": (
+            "foooooooooooooooooooo,\n"
+            "barrrrrrrrrrrrrrrrrrr,\n"
+            "bazzzzzzzzzzzzzzzzzzz,\n"
+        ),
+    } %}
+{% endmacro %}
diff --git a/tests/functional_tests/test_general_formatting.py b/tests/functional_tests/test_general_formatting.py
@@ -18,6 +18,7 @@
         "preformatted/008_reserved_names.sql",
         "preformatted/301_multiline_jinjafmt.sql",
         "preformatted/302_jinjafmt_multiline_str.sql",
+        "preformatted/303_jinjafmt_more_mutliline_str.sql",
         "preformatted/400_create_table.sql",
         "unformatted/100_select_case.sql",
         "unformatted/101_multiline.sql",