Fix a crash when # fmt: on is used on a different block level than # fmt: off (#3281)

Previously _Black_ produces invalid code because the `# fmt: on` is used on a different block level.

While _Black_ requires `# fmt: off` and `# fmt: on` to be used at the same block level, incorrect usage shouldn't cause crashes.

The formatting behavior this PR introduces is, the code below the initial `# fmt: off` block level will be turned off for formatting, when `# fmt: on` is used on a different level or there is no `# fmt: on`. This also matches the current behavior when `# fmt: off` is used at the top-level without a matching `# fmt: on`, it turns off formatting for everything below `# fmt: off`.

- Fixes #2567
- Fixes #3184
- Fixes #2985
- Fixes #2882
- Fixes #2232
- Fixes #2140
- Fixes #1817
- Fixes #569

CHANGES.md

@@ -10,6 +10,9 @@
 
 <!-- Changes that affect Black's stable style -->
 
+- Fix a crash when `# fmt: on` is used on a different block level than `# fmt: off`
+  (#3281)
+
 ### Preview style
 
 <!-- Changes that affect Black's preview style -->

docs/contributing/reference/reference_functions.rst

@@ -137,9 +137,9 @@ Utilities
 
 .. autofunction:: black.comments.is_fmt_on
 
-.. autofunction:: black.comments.contains_fmt_on_at_column
+.. autofunction:: black.comments.children_contains_fmt_on
 
-.. autofunction:: black.nodes.first_leaf_column
+.. autofunction:: black.nodes.first_leaf_of
 
 .. autofunction:: black.linegen.generate_trailers_to_omit
 

src/black/comments.py

@@ -14,11 +14,12 @@
     STANDALONE_COMMENT,
     WHITESPACE,
     container_of,
-    first_leaf_column,
+    first_leaf_of,
     preceding_leaf,
+    syms,
 )
 from blib2to3.pgen2 import token
-from blib2to3.pytree import Leaf, Node, type_repr
+from blib2to3.pytree import Leaf, Node
 
 # types
 LN = Union[Leaf, Node]
@@ -230,7 +231,7 @@ def generate_ignored_nodes(
             return
 
         # fix for fmt: on in children
-        if contains_fmt_on_at_column(container, leaf.column, preview=preview):
+        if children_contains_fmt_on(container, preview=preview):
             for child in container.children:
                 if isinstance(child, Leaf) and is_fmt_on(child, preview=preview):
                     if child.type in CLOSING_BRACKETS:
@@ -240,10 +241,14 @@ def generate_ignored_nodes(
                         # The alternative is to fail the formatting.
                         yield child
                     return
-                if contains_fmt_on_at_column(child, leaf.column, preview=preview):
+                if children_contains_fmt_on(child, preview=preview):
                     return
                 yield child
         else:
+            if container.type == token.DEDENT and container.next_sibling is None:
+                # This can happen when there is no matching `# fmt: on` comment at the
+                # same level as `# fmt: on`. We need to keep this DEDENT.
+                return
             yield container
             container = container.next_sibling
 
@@ -268,9 +273,7 @@ def _generate_ignored_nodes_from_fmt_skip(
         for sibling in siblings:
             yield sibling
     elif (
-        parent is not None
+        parent is not None and parent.type == syms.suite and leaf.type == token.NEWLINE
-        and type_repr(parent.type) == "suite"
-        and leaf.type == token.NEWLINE
     ):
         # The `# fmt: skip` is on the colon line of the if/while/def/class/...
         # statements. The ignored nodes should be previous siblings of the
@@ -278,7 +281,7 @@ def _generate_ignored_nodes_from_fmt_skip(
         leaf.prefix = ""
         ignored_nodes: List[LN] = []
         parent_sibling = parent.prev_sibling
-        while parent_sibling is not None and type_repr(parent_sibling.type) != "suite":
+        while parent_sibling is not None and parent_sibling.type != syms.suite:
             ignored_nodes.insert(0, parent_sibling)
             parent_sibling = parent_sibling.prev_sibling
         # Special case for `async_stmt` where the ASYNC token is on the
@@ -306,17 +309,12 @@ def is_fmt_on(container: LN, preview: bool) -> bool:
     return fmt_on
 
 
-def contains_fmt_on_at_column(container: LN, column: int, *, preview: bool) -> bool:
+def children_contains_fmt_on(container: LN, *, preview: bool) -> bool:
-    """Determine if children at a given column have formatting switched on."""
+    """Determine if children have formatting switched on."""
     for child in container.children:
-        if (
+        leaf = first_leaf_of(child)
-            isinstance(child, Node)
+        if leaf is not None and is_fmt_on(leaf, preview=preview):
-            and first_leaf_column(child) == column
+            return True
-            or isinstance(child, Leaf)
-            and child.column == column
-        ):
-            if is_fmt_on(child, preview=preview):
-                return True
 
     return False
 

src/black/nodes.py

@@ -502,12 +502,14 @@ def container_of(leaf: Leaf) -> LN:
     return container
 
 
-def first_leaf_column(node: Node) -> Optional[int]:
+def first_leaf_of(node: LN) -> Optional[Leaf]:
-    """Returns the column of the first leaf child of a node."""
+    """Returns the first leaf of the node tree."""
-    for child in node.children:
+    if isinstance(node, Leaf):
-        if isinstance(child, Leaf):
+        return node
-            return child.column
+    if node.children:
-    return None
+        return first_leaf_of(node.children[0])
+    else:
+        return None
 
 
 def is_arith_like(node: LN) -> bool:

tests/data/simple_cases/fmtonoff5.py

@@ -34,3 +34,125 @@ def test_func():
         return True
 
     return False
+
+
+# Regression test for https://github.com/psf/black/issues/2567.
+if True:
+    # fmt: off
+    for _ in range( 1 ):
+    # fmt: on
+        print ( "This won't be formatted" )
+    print ( "This won't be formatted either" )
+else:
+    print ( "This will be formatted" )
+
+
+# Regression test for https://github.com/psf/black/issues/3184.
+class A:
+    async def call(param):
+        if param:
+            # fmt: off
+            if param[0:4] in (
+                "ABCD", "EFGH"
+            )  :
+                # fmt: on
+                print ( "This won't be formatted" )
+
+            elif param[0:4] in ("ZZZZ",):
+                print ( "This won't be formatted either" )
+
+        print ( "This will be formatted" )
+
+
+# Regression test for https://github.com/psf/black/issues/2985
+class Named(t.Protocol):
+    # fmt: off
+    @property
+    def  this_wont_be_formatted ( self ) -> str: ...
+
+class Factory(t.Protocol):
+    def  this_will_be_formatted ( self, **kwargs ) -> Named: ...
+    # fmt: on
+
+
+# output
+
+
+# Regression test for https://github.com/psf/black/issues/3129.
+setup(
+    entry_points={
+        # fmt: off
+        "console_scripts": [
+            "foo-bar"
+            "=foo.bar.:main",
+        # fmt: on
+            ]  # Includes an formatted indentation.
+    },
+)
+
+
+# Regression test for https://github.com/psf/black/issues/2015.
+run(
+    # fmt: off
+    [
+        "ls",
+        "-la",
+    ]
+    # fmt: on
+    + path,
+    check=True,
+)
+
+
+# Regression test for https://github.com/psf/black/issues/3026.
+def test_func():
+    # yapf: disable
+    if  unformatted(  args  ):
+        return True
+    # yapf: enable
+    elif b:
+        return True
+
+    return False
+
+
+# Regression test for https://github.com/psf/black/issues/2567.
+if True:
+    # fmt: off
+    for _ in range( 1 ):
+    # fmt: on
+        print ( "This won't be formatted" )
+    print ( "This won't be formatted either" )
+else:
+    print("This will be formatted")
+
+
+# Regression test for https://github.com/psf/black/issues/3184.
+class A:
+    async def call(param):
+        if param:
+            # fmt: off
+            if param[0:4] in (
+                "ABCD", "EFGH"
+            )  :
+                # fmt: on
+                print ( "This won't be formatted" )
+
+            elif param[0:4] in ("ZZZZ",):
+                print ( "This won't be formatted either" )
+
+        print("This will be formatted")
+
+
+# Regression test for https://github.com/psf/black/issues/2985
+class Named(t.Protocol):
+    # fmt: off
+    @property
+    def  this_wont_be_formatted ( self ) -> str: ...
+
+
+class Factory(t.Protocol):
+    def this_will_be_formatted(self, **kwargs) -> Named:
+        ...
+
+    # fmt: on