fix: fix issue with whitespace and data tokens

tconbeer · tconbeer · Dec 12, 2022 · Dec 9, 2022 · Dec 9, 2022 · Dec 11, 2022
commit 4f9a91758f2c6e967ac51f4ee8d2d433f4c79533
diff --git a/src/sqlfmt/node_manager.py b/src/sqlfmt/node_manager.py
@@ -1,5 +1,5 @@
 import re
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 from sqlfmt.exception import SqlfmtBracketError
 from sqlfmt.line import Line
@@ -23,54 +23,15 @@ def create_node(self, token: Token, previous_node: Optional[Node]) -> Node:
         lowercased if they are simple names, keywords, or statements.
         """
 
-        if previous_node is None:
-            open_brackets = []
-            open_jinja_blocks = []
-        else:
-            open_brackets = previous_node.open_brackets.copy()
-            open_jinja_blocks = previous_node.open_jinja_blocks.copy()
-
-            # add the previous node to the list of open brackets or jinja blocks
-            if previous_node.is_unterm_keyword or previous_node.is_opening_bracket:
-                open_brackets.append(previous_node)
-            elif previous_node.is_opening_jinja_block:
-                open_jinja_blocks.append(previous_node)
-
-        # if the token should reduce the depth of the node, pop
-        # the last item(s) off open_brackets or open_jinja_blocks
-        if token.type in (TokenType.UNTERM_KEYWORD, TokenType.SET_OPERATOR):
-            if open_brackets and open_brackets[-1].is_unterm_keyword:
-                _ = open_brackets.pop()
-        elif token.type in (TokenType.BRACKET_CLOSE, TokenType.STATEMENT_END):
-            try:
-                last_bracket = open_brackets.pop()
-                if last_bracket.is_unterm_keyword:
-                    last_bracket = open_brackets.pop()
-            except IndexError:
-                raise SqlfmtBracketError(
-                    f"Closing bracket '{token.token}' found at "
-                    f"{token.spos} before bracket was opened."
-                )
-            else:
-                self.raise_on_mismatched_bracket(token, last_bracket)
-        elif token.type is TokenType.JINJA_BLOCK_END:
-            try:
-                start_tag = open_jinja_blocks.pop()
-                self.raise_on_mismatched_jinja_tags(token, start_tag)
-            except IndexError:
-                raise SqlfmtBracketError(
-                    f"Closing bracket '{token.token}' found at "
-                    f"{token.spos} before bracket was opened."
-                )
-        # if we hit a semicolon, reset open_brackets, since we're
-        # about to start a new query
-        elif token.type is TokenType.SEMICOLON:
-            open_brackets = []
-
-        prev_token, extra_whitespace = get_previous_token(previous_node)
-        prefix = self.whitespace(token, prev_token, extra_whitespace)
-        value = self.standardize_value(token)
+        open_brackets, open_jinja_blocks = self.open_brackets(token, previous_node)
         formatting_disabled = self.disable_formatting(token, previous_node)
+        if formatting_disabled:
+            prefix = token.prefix
+            value = token.token
+        else:
+            prev_token, extra_whitespace = get_previous_token(previous_node)
+            prefix = self.whitespace(token, prev_token, extra_whitespace)
+            value = self.standardize_value(token)
 
         return Node(
             token=token,
@@ -109,6 +70,10 @@ def raise_on_mismatched_bracket(self, token: Token, last_bracket: Node) -> None:
             )
 
     def raise_on_mismatched_jinja_tags(self, token: Token, start_tag: Node) -> None:
+        """
+        Compare the value of token to the start_tag to determine whether token
+        closes start_tag
+        """
         try:
             if any(s in token.token.lower() for s in ["endif", "else", "elif"]):
                 if not any(s in start_tag.value for s in ["if", "else"]):
@@ -125,6 +90,62 @@ def raise_on_mismatched_jinja_tags(self, token: Token, start_tag: Node) -> None:
                 f"{start_tag.token.spos}."
             )
 
+    def open_brackets(
+        self, token: Token, previous_node: Optional[Node]
+    ) -> Tuple[List[Node], List[Node]]:
+        """
+        Uses the previous_node and the contents of the current token
+        to compute the depth of the new node.
+
+        Returns two lists, for open_brackets and open_jinja_blocks
+        """
+
+        if previous_node is None:
+            open_brackets = []
+            open_jinja_blocks = []
+        else:
+            open_brackets = previous_node.open_brackets.copy()
+            open_jinja_blocks = previous_node.open_jinja_blocks.copy()
+
+            # add the previous node to the list of open brackets or jinja blocks
+            if previous_node.is_unterm_keyword or previous_node.is_opening_bracket:
+                open_brackets.append(previous_node)
+            elif previous_node.is_opening_jinja_block:
+                open_jinja_blocks.append(previous_node)
+
+        # if the token should reduce the depth of the node, pop
+        # the last item(s) off open_brackets or open_jinja_blocks
+        if token.type in (TokenType.UNTERM_KEYWORD, TokenType.SET_OPERATOR):
+            if open_brackets and open_brackets[-1].is_unterm_keyword:
+                _ = open_brackets.pop()
+        elif token.type in (TokenType.BRACKET_CLOSE, TokenType.STATEMENT_END):
+            try:
+                last_bracket = open_brackets.pop()
+                if last_bracket.is_unterm_keyword:
+                    last_bracket = open_brackets.pop()
+            except IndexError:
+                raise SqlfmtBracketError(
+                    f"Closing bracket '{token.token}' found at "
+                    f"{token.spos} before bracket was opened."
+                )
+            else:
+                self.raise_on_mismatched_bracket(token, last_bracket)
+        elif token.type is TokenType.JINJA_BLOCK_END:
+            try:
+                start_tag = open_jinja_blocks.pop()
+                self.raise_on_mismatched_jinja_tags(token, start_tag)
+            except IndexError:
+                raise SqlfmtBracketError(
+                    f"Closing bracket '{token.token}' found at "
+                    f"{token.spos} before bracket was opened."
+                )
+        # if we hit a semicolon, reset open_brackets, since we're
+        # about to start a new query
+        elif token.type is TokenType.SEMICOLON:
+            open_brackets = []
+
+        return open_brackets, open_jinja_blocks
+
     def whitespace(
         self,
         token: Token,

diff --git a/src/sqlfmt/splitter.py b/src/sqlfmt/splitter.py
@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Iterator
+from typing import Iterator, Tuple
 
 from sqlfmt.line import Line
 from sqlfmt.node import Node
@@ -20,17 +20,22 @@ def maybe_split(self, line: Line) -> Iterator[Line]:
             yield line
             return
 
-        split_after = False
+        always_split_after = never_split_after = False
         for i, node in enumerate(line.nodes):
             if node.is_newline:
                 # can't split just before a newline
                 yield line
                 break
-            elif i > 0 and (split_after or self.maybe_split_before(node)):
+            elif (
+                i > 0
+                and not never_split_after
+                and not node.formatting_disabled
+                and (always_split_after or self.maybe_split_before(node))
+            ):
                 yield from self.split_at_index(line, i)
                 break
 
-            split_after = self.maybe_split_after(node)
+            always_split_after, never_split_after = self.maybe_split_after(node)
 
     def maybe_split_before(self, node: Node) -> bool:
         """
@@ -73,9 +78,12 @@ def maybe_split_between_brackets(self, node: Node) -> bool:
             and node.previous_node.is_closing_bracket
         )
 
-    def maybe_split_after(self, node: Node) -> bool:
+    def maybe_split_after(self, node: Node) -> Tuple[bool, bool]:
         """
-        Return True if we should split after node
+        Return True, False if we should always split after node
+        Retrun False, True if we should never split after node
+        Return False, False if splitting after should depend on the
+        contents of the next node
         """
         if (
             # always split after any comma that doesn't end a line
@@ -87,9 +95,11 @@ def maybe_split_after(self, node: Node) -> bool:
             # always split after a token that divides queries
             or node.divides_queries
         ):
-            return True
+            return True, False
+        elif node.formatting_disabled:
+            return False, True
         else:
-            return False
+            return False, False
 
     def split_at_index(self, line: Line, index: int) -> Iterator[Line]:
         """

diff --git a/src/sqlfmt_primer/primer.py b/src/sqlfmt_primer/primer.py
@@ -30,7 +30,7 @@ def get_projects() -> List[SQLProject]:
         SQLProject(
             name="gitlab",
             git_url="https://github.com/tconbeer/gitlab-analytics-sqlfmt.git",
-            git_ref="90ffb23",  # sqlfmt a7ed980
+            git_ref="091e15b",  # sqlfmt 7be6ac5
             expected_changed=3,
             expected_unchanged=2414,
             expected_errored=0,

diff --git a/tests/unit_tests/test_splitter.py b/tests/unit_tests/test_splitter.py
@@ -325,3 +325,15 @@ def test_split_between_brackets(
         "    )\n",
     ]
     assert actual_result == expected_result
+
+
+def test_split_around_data(splitter: LineSplitter, default_analyzer: Analyzer) -> None:
+    source_string = "{% set foo %}//my one-line //data{% endset %}\n"
+    raw_query = default_analyzer.parse_query(source_string)
+
+    split_lines: List[Line] = []
+    for raw_line in raw_query.lines:
+        split_lines.extend(splitter.maybe_split(raw_line))
+
+    actual_result = [str(line) for line in split_lines]
+    assert actual_result == [source_string]