Skip to content

Commit 032ce0f

Browse files
keraionkeraion
authored
ST04: Retain comments when flattening CASE (#5753)
1 parent 081a526 commit 032ce0f

2 files changed

Lines changed: 204 additions & 15 deletions

File tree

  • src/sqlfluff/rules/structure
  • test/fixtures/rules/std_rule_cases

‎src/sqlfluff/rules/structure/ST04.py‎

Lines changed: 120 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
"""Implementation of Rule ST04."""
22

3-
from sqlfluff.core.parser import NewlineSegment, WhitespaceSegment
3+
from typing import List
4+
5+
from sqlfluff.core.parser import BaseSegment, Indent, NewlineSegment, WhitespaceSegment
46
from sqlfluff.core.rules import BaseRule, LintFix, LintResult, RuleContext
57
from sqlfluff.core.rules.crawlers import SegmentSeekerCrawler
6-
from sqlfluff.utils.functional import FunctionalContext, sp
8+
from sqlfluff.utils.functional import FunctionalContext, Segments, sp
79
from sqlfluff.utils.reflow.reindent import construct_single_indent
810

911

@@ -53,6 +55,7 @@ def _eval(self, context: RuleContext) -> LintResult:
5355
assert segment.select(sp.is_type("case_expression"))
5456
case1_children = segment.children()
5557
case1_first_case = case1_children.first(sp.is_keyword("CASE")).get()
58+
assert case1_first_case
5659
case1_first_when = case1_children.first(
5760
sp.is_type("when_clause", "else_clause")
5861
).get()
@@ -104,36 +107,138 @@ def _eval(self, context: RuleContext) -> LintResult:
104107
case1_to_delete = case1_children.select(
105108
start_seg=case1_last_when, stop_seg=case1_else_clause_seg
106109
)
110+
# Restore any comments that were deleted
111+
after_last_comment_index = (
112+
case1_to_delete.find(case1_to_delete.last(sp.is_comment()).get()) + 1
113+
)
114+
case1_comments_to_restore = case1_to_delete.select(
115+
stop_seg=case1_to_delete.get(after_last_comment_index)
116+
)
117+
after_else_comment = case1_else_clause.children().select(
118+
select_if=sp.is_type("newline", "comment", "whitespace"),
119+
stop_seg=case1_else_expressions.get(),
120+
)
107121

108122
# Delete the nested "CASE" expression.
109-
fixes = case1_to_delete.apply(lambda seg: LintFix.delete(seg))
123+
fixes = case1_to_delete.apply(LintFix.delete)
110124

111125
tab_space_size: int = context.config.get("tab_space_size", ["indentation"])
112126
indent_unit: str = context.config.get("indent_unit", ["indentation"])
113127

114128
# Determine the indentation to use when we move the nested "WHEN"
115129
# and "ELSE" clauses, based on the indentation of case1_last_when.
116130
# If no whitespace segments found, use default indent.
117-
indent = (
118-
case1_children.select(stop_seg=case1_last_when)
119-
.reversed()
120-
.first(sp.is_type("whitespace"))
131+
when_indent_str = self._get_indentation(
132+
case1_children, case1_last_when, tab_space_size, indent_unit
121133
)
122-
indent_str = (
123-
"".join(seg.raw for seg in indent)
124-
if indent
125-
else construct_single_indent(indent_unit, tab_space_size)
134+
# Again determine indentation, but matching the "CASE"/"END" level.
135+
end_indent_str = self._get_indentation(
136+
case1_children, case1_first_case, tab_space_size, indent_unit
126137
)
127138

128139
# Move the nested "when" and "else" clauses after the last outer
129140
# "when".
130-
nested_clauses = case2.children(sp.is_type("when_clause", "else_clause"))
131-
create_after_last_when = nested_clauses.apply(
132-
lambda seg: [NewlineSegment(), WhitespaceSegment(indent_str), seg]
141+
nested_clauses = case2.children(
142+
sp.is_type("when_clause", "else_clause", "newline", "comment", "whitespace")
133143
)
134-
segments = [item for sublist in create_after_last_when for item in sublist]
144+
145+
# Rebuild the nested case statement.
146+
# Any comments after the last outer "WHEN" that were deleted
147+
segments = list(case1_comments_to_restore)
148+
# Any comments between the "ELSE" and nested "CASE"
149+
segments += self._rebuild_spacing(when_indent_str, after_else_comment)
150+
# The nested "WHEN", "ELSE" or "comments", with logical spacing
151+
segments += self._rebuild_spacing(when_indent_str, nested_clauses)
135152
fixes.append(LintFix.create_after(case1_last_when, segments, source=segments))
136153

137154
# Delete the outer "else" clause.
138155
fixes.append(LintFix.delete(case1_else_clause_seg))
156+
# Add spacing for any comments that may exist after the nested `END`
157+
# but only on that same line.
158+
fixes += self._nested_end_trailing_comment(
159+
case1_children, case1_else_clause_seg, end_indent_str
160+
)
139161
return LintResult(case2[0], fixes=fixes)
162+
163+
def _get_indentation(
164+
self,
165+
parent_segments: Segments,
166+
segment: BaseSegment,
167+
tab_space_size: int,
168+
indent_unit: str,
169+
) -> str:
170+
"""Calculate the indentation level for rebuilding nested struct.
171+
172+
This is only a best attempt as the input may not be equally indented. The layout
173+
rules, if run, would resolve this.
174+
"""
175+
leading_whitespace = (
176+
parent_segments.select(stop_seg=segment)
177+
.reversed()
178+
.first(sp.is_type("whitespace"))
179+
)
180+
seg_indent = parent_segments.select(stop_seg=segment).last(sp.is_type("indent"))
181+
indent_level = 1
182+
if (
183+
seg_indent
184+
and (segment_indent := seg_indent.get())
185+
and isinstance(segment_indent, Indent)
186+
):
187+
indent_level = segment_indent.indent_val + 1
188+
indent_str = (
189+
"".join(seg.raw for seg in leading_whitespace)
190+
if leading_whitespace
191+
and (whitespace_seg := leading_whitespace.get())
192+
and len(whitespace_seg.raw) > 1
193+
else construct_single_indent(indent_unit, tab_space_size) * indent_level
194+
)
195+
196+
return indent_str
197+
198+
def _nested_end_trailing_comment(
199+
self,
200+
case1_children: Segments,
201+
case1_else_clause_seg: BaseSegment,
202+
end_indent_str: str,
203+
) -> List[LintFix]:
204+
"""Prepend newline spacing to comments on the final nested `END` line."""
205+
trailing_end = case1_children.select(
206+
start_seg=case1_else_clause_seg,
207+
loop_while=sp.not_(sp.is_type("newline")),
208+
)
209+
fixes = trailing_end.select(
210+
sp.is_whitespace(), loop_while=sp.not_(sp.is_comment())
211+
).apply(LintFix.delete)
212+
first_comment = trailing_end.first(sp.is_comment()).get()
213+
if first_comment:
214+
segments = [NewlineSegment(), WhitespaceSegment(end_indent_str)]
215+
fixes.append(LintFix.create_before(first_comment, segments, segments))
216+
return fixes
217+
218+
def _rebuild_spacing(
219+
self, indent_str: str, nested_clauses: Segments
220+
) -> List[BaseSegment]:
221+
buff = []
222+
# If the first segment is a comment, add a newline
223+
prior_newline = nested_clauses.first(sp.not_(sp.is_whitespace())).any(
224+
sp.is_comment()
225+
)
226+
prior_whitespace = ""
227+
for seg in nested_clauses:
228+
if seg.is_type("when_clause", "else_clause") or (
229+
prior_newline and seg.is_comment
230+
):
231+
buff += [NewlineSegment(), WhitespaceSegment(indent_str), seg]
232+
prior_newline = False
233+
prior_whitespace = ""
234+
elif seg.is_type("newline"):
235+
prior_newline = True
236+
prior_whitespace = ""
237+
elif not prior_newline and seg.is_comment:
238+
buff += [WhitespaceSegment(prior_whitespace), seg]
239+
prior_newline = False
240+
prior_whitespace = ""
241+
elif seg.is_whitespace:
242+
# Don't reset newline
243+
prior_whitespace = seg.raw
244+
return buff

‎test/fixtures/rules/std_rule_cases/ST04.yml‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,3 +284,87 @@ test_fail_nested_same_case:
284284
ELSE 'other'
285285
END
286286
FROM tab_a;
287+
288+
test_fail_retain_comments:
289+
fail_str: |
290+
SELECT
291+
CASE
292+
WHEN FALSE
293+
THEN "value1" -- a comment
294+
ELSE
295+
CASE
296+
-- another comment
297+
WHEN TRUE -- and here
298+
THEN "value2" -- but also here
299+
END
300+
END
301+
FROM table;
302+
fix_str: |
303+
SELECT
304+
CASE
305+
WHEN FALSE
306+
THEN "value1" -- a comment
307+
-- another comment
308+
WHEN TRUE -- and here
309+
THEN "value2" -- but also here
310+
END
311+
FROM table;
312+
313+
test_fail_retain_comments_after_end:
314+
fail_str: |
315+
SELECT
316+
CASE -- no spaces here
317+
WHEN FALSE
318+
THEN "value1" -- a comment
319+
ELSE
320+
CASE -- after case
321+
-- another comment
322+
/* before the when */ WHEN TRUE -- and here
323+
THEN "value2" -- but also here
324+
END /* after the end */ /* but wait there's more! */
325+
-- but here too
326+
END
327+
FROM table;
328+
fix_str: |
329+
SELECT
330+
CASE -- no spaces here
331+
WHEN FALSE
332+
THEN "value1" -- a comment
333+
-- after case
334+
-- another comment
335+
/* before the when */
336+
WHEN TRUE -- and here
337+
THEN "value2" -- but also here
338+
/* after the end */ /* but wait there's more! */
339+
-- but here too
340+
END
341+
FROM table;
342+
343+
test_fail_retain_comments_after_else:
344+
fail_str: |
345+
SELECT
346+
CASE
347+
WHEN FALSE
348+
THEN "value1" -- a comment
349+
/* before else*/ ELSE --after else
350+
/*before case*/ CASE -- else case
351+
-- another comment
352+
WHEN TRUE -- and here
353+
THEN "value2" -- but also here
354+
END
355+
END
356+
FROM table;
357+
fix_str: |
358+
SELECT
359+
CASE
360+
WHEN FALSE
361+
THEN "value1" -- a comment
362+
/* before else*/
363+
--after else
364+
/*before case*/
365+
-- else case
366+
-- another comment
367+
WHEN TRUE -- and here
368+
THEN "value2" -- but also here
369+
END
370+
FROM table;

0 commit comments

Comments
 (0)