Fix crashes when fmt directives are indented (#4856)

* Use `_contains_fmt_directive` everywhere

Signed-off-by: cobalt <[email protected]>

* Remove underscore prefix

Signed-off-by: cobalt <[email protected]>

* Add tests

Signed-off-by: cobalt <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update changelog

Signed-off-by: cobalt <[email protected]>

* Remove duplicate tests

* fix `visit_STANDALONE_COMMENT`'s `is_fmt_off_block` check

Signed-off-by: cobalt <[email protected]>

* Prevent wrapping when comments are in arrays

Signed-off-by: cobalt <[email protected]>

---------

Signed-off-by: cobalt <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: cobaltt7

CHANGES.md

@@ -17,6 +17,7 @@
 
 - Fix bug where comments preceding `# fmt: off`/`# fmt: on` blocks were incorrectly
   removed, particularly affecting Jupytext's `# %% [markdown]` comments (#4845)
+- Fix possible crash when `fmt: ` directives aren't on the top level (#4856)
 
 ### Preview style
 

src/black/comments.py

@@ -164,7 +164,7 @@ def make_comment(content: str, mode: Mode) -> str:
         return "#"
 
     # Preserve comments with fmt directives exactly as-is
-    if content.startswith("#") and _contains_fmt_directive(content):
+    if content.startswith("#") and contains_fmt_directive(content):
         return content
 
     if content[0] == "#":
@@ -205,8 +205,8 @@ def _should_process_fmt_comment(
 
     Returns (should_process, is_fmt_off, is_fmt_skip).
     """
-    is_fmt_off = _contains_fmt_directive(comment.value, FMT_OFF)
-    is_fmt_skip = _contains_fmt_directive(comment.value, FMT_SKIP)
+    is_fmt_off = contains_fmt_directive(comment.value, FMT_OFF)
+    is_fmt_skip = contains_fmt_directive(comment.value, FMT_SKIP)
 
     if not is_fmt_off and not is_fmt_skip:
         return False, False, False
@@ -258,9 +258,13 @@ def _handle_comment_only_fmt_block(
     fmt_off_idx = None
     fmt_on_idx = None
     for idx, c in enumerate(all_comments):
-        if fmt_off_idx is None and c.value in FMT_OFF:
+        if fmt_off_idx is None and contains_fmt_directive(c.value, FMT_OFF):
             fmt_off_idx = idx
-        if fmt_off_idx is not None and idx > fmt_off_idx and c.value in FMT_ON:
+        if (
+            fmt_off_idx is not None
+            and idx > fmt_off_idx
+            and contains_fmt_directive(c.value, FMT_ON)
+        ):
             fmt_on_idx = idx
             break
 
@@ -401,7 +405,7 @@ def _handle_regular_fmt_block(
     parent = first.parent
     prefix = first.prefix
 
-    if comment.value in FMT_OFF:
+    if contains_fmt_directive(comment.value, FMT_OFF):
         first.prefix = prefix[comment.consumed :]
     if is_fmt_skip:
         first.prefix = ""
@@ -412,7 +416,7 @@ def _handle_regular_fmt_block(
     hidden_value = "".join(str(n) for n in ignored_nodes)
     comment_lineno = leaf.lineno - comment.newlines
 
-    if comment.value in FMT_OFF:
+    if contains_fmt_directive(comment.value, FMT_OFF):
         fmt_off_prefix = ""
         if len(lines) > 0 and not any(
             line[0] <= comment_lineno <= line[1] for line in lines
@@ -460,7 +464,7 @@ def generate_ignored_nodes(
     If comment is skip, returns leaf only.
     Stops at the end of the block.
     """
-    if _contains_fmt_directive(comment.value, FMT_SKIP):
+    if contains_fmt_directive(comment.value, FMT_SKIP):
         yield from _generate_ignored_nodes_from_fmt_skip(leaf, comment, mode)
         return
     container: LN | None = container_of(leaf)
@@ -717,9 +721,9 @@ def is_fmt_on(container: LN, mode: Mode) -> bool:
     """
     fmt_on = False
     for comment in list_comments(container.prefix, is_endmarker=False, mode=mode):
-        if comment.value in FMT_ON:
+        if contains_fmt_directive(comment.value, FMT_ON):
             fmt_on = True
-        elif comment.value in FMT_OFF:
+        elif contains_fmt_directive(comment.value, FMT_OFF):
             fmt_on = False
     return fmt_on
 
@@ -748,7 +752,7 @@ def contains_pragma_comment(comment_list: list[Leaf]) -> bool:
     return False
 
 
-def _contains_fmt_directive(
+def contains_fmt_directive(
     comment_line: str, directives: set[str] = FMT_OFF | FMT_ON | FMT_SKIP
 ) -> bool:
     """

src/black/linegen.py

@@ -20,7 +20,7 @@
 from black.comments import (
     FMT_OFF,
     FMT_ON,
-    _contains_fmt_directive,
+    contains_fmt_directive,
     generate_comments,
     list_comments,
 )
@@ -387,7 +387,8 @@ def visit_ENDMARKER(self, leaf: Leaf) -> Iterator[Line]:
         yield from self.line()
 
     def visit_STANDALONE_COMMENT(self, leaf: Leaf) -> Iterator[Line]:
-        if not self.current_line.bracket_tracker.any_open_brackets():
+        any_open_brackets = self.current_line.bracket_tracker.any_open_brackets()
+        if not any_open_brackets:
             yield from self.line()
         # STANDALONE_COMMENT nodes created by our special handling in
         # normalize_fmt_off for comment-only blocks have fmt:off as the first
@@ -398,18 +399,11 @@ def visit_STANDALONE_COMMENT(self, leaf: Leaf) -> Iterator[Line]:
         # visit_default.
         value = leaf.value
         lines = value.splitlines()
-        if len(lines) >= 2:
-            # Check if first line (after stripping whitespace) is exactly a
-            # fmt:off directive
-            first_line = lines[0].lstrip()
-            first_is_fmt_off = first_line in FMT_OFF
-            # Check if last line (after stripping whitespace) is exactly a
-            # fmt:on directive
-            last_line = lines[-1].lstrip()
-            last_is_fmt_on = last_line in FMT_ON
-            is_fmt_off_block = first_is_fmt_off and last_is_fmt_on
-        else:
-            is_fmt_off_block = False
+        is_fmt_off_block = (
+            len(lines) >= 2
+            and contains_fmt_directive(lines[0], FMT_OFF)
+            and contains_fmt_directive(lines[-1], FMT_ON)
+        )
         if is_fmt_off_block:
             # This is a fmt:off/on block from normalize_fmt_off - we still need
             # to process any prefix comments (like markdown comments) but append
@@ -418,7 +412,7 @@ def visit_STANDALONE_COMMENT(self, leaf: Leaf) -> Iterator[Line]:
             # Only process prefix comments if there actually is a prefix with comments
             if leaf.prefix and any(
                 line.strip().startswith("#")
-                and not _contains_fmt_directive(line.strip())
+                and not contains_fmt_directive(line.strip())
                 for line in leaf.prefix.split("\n")
             ):
                 for comment in generate_comments(leaf, mode=self.mode):
@@ -429,7 +423,8 @@ def visit_STANDALONE_COMMENT(self, leaf: Leaf) -> Iterator[Line]:
                 leaf.prefix = ""
 
             self.current_line.append(leaf)
-            yield from self.line()
+            if not any_open_brackets:
+                yield from self.line()
         else:
             # Normal standalone comment - process through visit_default
             yield from self.visit_default(leaf)
@@ -1484,7 +1479,7 @@ def normalize_invisible_parens(
     existing visible parentheses for other tuples and generator expressions.
     """
     for pc in list_comments(node.prefix, is_endmarker=False, mode=mode):
-        if pc.value in FMT_OFF:
+        if contains_fmt_directive(pc.value, FMT_OFF):
             # This `node` has a prefix with `# fmt: off`, don't mess with parens.
             return
 

tests/data/cases/fmtskip11.py

@@ -4,3 +4,85 @@ def foo():
 
 # comment 1  # fmt: skip
 # comment 2
+
+[
+    (1, 2),
+    # # fmt: off
+    # (3,
+    #    4),
+    # # fmt: on
+    (5, 6),
+]
+
+[
+    (1, 2),
+    # # fmt: off
+    # (3,
+    #    4),
+    # fmt: on
+    (5, 6),
+]
+
+
+[
+    (1, 2),
+    # fmt: off
+    # (3,
+    #    4),
+    # # fmt: on
+    (5, 6),
+]
+
+
+[
+    (1, 2),
+    # fmt: off
+    # (3,
+    #    4),
+    # fmt: on
+    (5, 6),
+]
+
+[
+    (1, 2),
+    # # fmt: off
+    (3,
+       4),
+    # # fmt: on
+    (5, 6),
+]
+
+[
+    (1, 2),
+    # # fmt: off
+    (3,
+       4),
+    # fmt: on
+    (5, 6),
+]
+
+
+[
+    (1, 2),
+    # fmt: off
+    (3,
+       4),
+    # # fmt: on
+    (5, 6),
+]
+
+
+[
+    (1, 2),
+    # fmt: off
+    (3,
+       4),
+    # fmt: on
+    (5, 6),
+]
+
+
+if False:
+    # fmt: off # some other comment
+    pass
+