[563] Fix standalone comments inside complex blocks crashing Black

CHANGES.md

@@ -13,6 +13,7 @@
 
 - Fix crash on formatting bytes strings that look like docstrings (#4003)
 - Fix crash when whitespace followed a backslash before newline in a docstring (#4008)
+- Fix standalone comments inside complex blocks crashing Black (#4016)
 
 - Fix crash on formatting code like `await (a ** b)` (#3994)
 

src/black/brackets.py

@@ -127,6 +127,13 @@ def mark(self, leaf: Leaf) -> None:
         self.maybe_increment_lambda_arguments(leaf)
         self.maybe_increment_for_loop_variable(leaf)
 
+    def any_open_for_or_lambda(self) -> bool:
+        """Return True if there is an open for or lambda expression on the line.
+
+        See maybe_increment_for_loop_variable and maybe_increment_lambda_arguments
+        for details."""
+        return bool(self._for_loop_depths or self._lambda_argument_depths)
+
     def any_open_brackets(self) -> bool:
         """Return True if there is an yet unmatched open bracket on the line."""
         return bool(self.bracket_match)

src/black/linegen.py

@@ -861,8 +861,6 @@ def _maybe_split_omitting_optional_parens(
         # it's not an import (optional parens are the only thing we can split on
         # in this case; attempting a split without them is a waste of time)
         and not line.is_import
-        # there are no standalone comments in the body
-        and not rhs.body.contains_standalone_comments(0)
         # and we can actually remove the parens
         and can_omit_invisible_parens(rhs, mode.line_length)
     ):
@@ -1181,7 +1179,7 @@ def standalone_comment_split(
     line: Line, features: Collection[Feature], mode: Mode
 ) -> Iterator[Line]:
     """Split standalone comments from the rest of the line."""
-    if not line.contains_standalone_comments(0):
+    if not line.contains_standalone_comments():
         raise CannotSplit("Line does not have any standalone comments")
 
     current_line = Line(

src/black/lines.py

@@ -1,6 +1,5 @@
 import itertools
 import math
-import sys
 from dataclasses import dataclass, field
 from typing import (
     Callable,
@@ -103,7 +102,10 @@ def append_safe(self, leaf: Leaf, preformatted: bool = False) -> None:
         Raises ValueError when any `leaf` is appended after a standalone comment
         or when a standalone comment is not the first leaf on the line.
         """
-        if self.bracket_tracker.depth == 0:
+        if (
+            self.bracket_tracker.depth == 0
+            or self.bracket_tracker.any_open_for_or_lambda()
+        ):
             if self.is_comment:
                 raise ValueError("cannot append to standalone comments")
 
@@ -233,10 +235,10 @@ def is_fmt_pass_converted(
             leaf.fmt_pass_converted_first_leaf
         )
 
-    def contains_standalone_comments(self, depth_limit: int = sys.maxsize) -> bool:
+    def contains_standalone_comments(self) -> bool:
         """If so, needs to be split before emitting."""
         for leaf in self.leaves:
-            if leaf.type == STANDALONE_COMMENT and leaf.bracket_depth <= depth_limit:
+            if leaf.type == STANDALONE_COMMENT:
                 return True
 
         return False
@@ -982,6 +984,23 @@ def can_omit_invisible_parens(
     are too long.
     """
     line = rhs.body
+
+    # We need optional parens in order to split standalone comments to their own lines
+    # if there are no nested parens around the standalone comments
+    closing_bracket: Optional[Leaf] = None
+    for leaf in reversed(line.leaves):
+        if closing_bracket and leaf is closing_bracket.opening_bracket:
+            closing_bracket = None
+        if leaf.type == STANDALONE_COMMENT and not closing_bracket:
+            return False
+        if (
+            not closing_bracket
+            and leaf.type in CLOSING_BRACKETS
+            and leaf.opening_bracket in line.leaves
+            and leaf.value
+        ):
+            closing_bracket = leaf
+
     bt = line.bracket_tracker
     if not bt.delimiters:
         # Without delimiters the optional parentheses are useless.

tests/data/cases/comments_in_blocks.py

@@ -0,0 +1,111 @@
+# Test cases from:
+# - https://github.com/psf/black/issues/1798
+# - https://github.com/psf/black/issues/1499
+# - https://github.com/psf/black/issues/1211
+# - https://github.com/psf/black/issues/563
+
+(
+    lambda
+    # a comment
+    : None
+)
+
+(
+    lambda:
+    # b comment
+    None
+)
+
+(
+    lambda
+    # a comment
+    :
+    # b comment
+    None
+)
+
+[
+    x
+    # Let's do this
+    for
+    # OK?
+    x
+    # Some comment
+    # And another
+    in
+    # One more
+    y
+]
+
+return [
+    (offers[offer_index], 1.0)
+    for offer_index, _
+    # avoid returning any offers that don't match the grammar so
+    # that the return values here are consistent with what would be
+    # returned in AcceptValidHeader
+    in self._parse_and_normalize_offers(offers)
+]
+
+from foo import (
+    bar,
+    # qux
+)
+
+
+def convert(collection):
+    # replace all variables by integers
+    replacement_dict = {
+        variable: f"{index}"
+        for index, variable
+        # 0 is reserved as line terminator
+        in enumerate(collection.variables(), start=1)
+    }
+
+
+{
+    i: i
+    for i
+    # a comment
+    in range(5)
+}
+
+
+def get_subtree_proof_nodes(
+    chunk_index_groups: Sequence[Tuple[int, ...], ...],
+) -> Tuple[int, ...]:
+    subtree_node_paths = (
+        # We take a candidate element from each group and shift it to
+        # remove the bits that are not common to other group members, then
+        # we convert it to a tree path that all elements from this group
+        # have in common.
+        chunk_index
+        for chunk_index, bits_to_truncate
+        # Each group will contain an even "power-of-two" number of# elements.
+        # This tells us how many tailing bits each element has# which need to
+        # be truncated to get the group's common prefix.
+        in ((group[0], (len(group) - 1).bit_length()) for group in chunk_index_groups)
+    )
+    return subtree_node_paths
+
+
+if (
+    # comment1
+    a
+    # comment2
+    or (
+        # comment3
+        (
+            # comment4
+            b
+        )
+        # comment5
+        and
+        # comment6
+        c
+        or (
+            # comment7
+            d
+        )
+    )
+):
+    print("Foo")

tests/data/cases/preview_hug_parens_with_braces_and_square_brackets.py

@@ -152,6 +152,16 @@ def foo_square_brackets(request):
 
 foo(**{x: y for x, y in enumerate(["long long long long line","long long long long line"])})
 
+for foo in ["a", "b"]:
+    output.extend([
+        individual
+        for
+        # Foobar
+        container in xs_by_y[foo]
+        # Foobar
+        for individual in container["nested"]
+    ])
+
 # output
 def foo_brackets(request):
     return JsonResponse({
@@ -323,3 +333,13 @@ def foo_square_brackets(request):
 foo(**{
     x: y for x, y in enumerate(["long long long long line", "long long long long line"])
 })
+
+for foo in ["a", "b"]:
+    output.extend([
+        individual
+        for
+        # Foobar
+        container in xs_by_y[foo]
+        # Foobar
+        for individual in container["nested"]
+    ])