Repository: psf/black

Issue:
Invalid code after indented `# fmt: off`
Input:
```py3
def x():
    # fmt: off
    return

def y():
    return
```

Invalid code produced:
```py3
def x():
    # fmt: off
    return

    def y():
        return
```


Discussion:
Is this supposed to work?

It leaves the file unchanged, but I would expect something like this

`in.py`
```python
# fmt: off
def x():
    return pow(
2,3)
    # fmt: on
    return pow(
2,3
)
```

`out.py`
```python
# fmt: off
def x():
    return pow(
2,3)
    # fmt: on
    return pow(2,3)
```

---

```
~ ) cat test.py 
# fmt: off
def x():
    return pow(
2,3)
    # fmt: on
    return pow(
2,3
)
~ ) black test.py
All done! ✨ 🍰 ✨
1 file left unchanged.
~ ) cat test.py 
# fmt: off
def x():
    return pow(
2,3)
    # fmt: on
    return pow(
2,3
)
```
Another interesting case: 

```python3
import asyncio

import aiohttp
from _pytest.monkeypatch import MonkeyPatch

from app import test_utils


async def test_ping_1(
    aiohttp_client: aiohttp.test_utils.TestClient,
    loop: asyncio.AbstractEventLoop,
    monkeypatch: MonkeyPatch,
):
  test_cases = (
    # test_input, expected
    ({"http_method": "HEAD"}, {"status": 204}),
    ({"http_method": "GET"}, {"status": 204}),
    ({"http_method": "POST"}, {"status": 405}),
  )

  # fmt: off
  async with \
      test_utils.ctx_test_config(monkeypatch=monkeypatch), \
      test_utils.ctx_mqtt_broker(loop=loop, monkeypatch=monkeypatch), \
      test_utils.ctx_app(loop=loop) as app \
  :
    # fmt: on
    client = await aiohttp_client(app)

    for (test_input, expected) in test_cases:
      resp = await client.request(test_input["http_method"], "/ping")
      assert resp.status == expected["status"]


# N.B. exact same function than `test_ping_1`, just there to
# illustrate the fact that Black "thinks" that `test_ping_2` is inside `test_ping_1`
async def test_ping_2(
    aiohttp_client: aiohttp.test_utils.TestClient,
    loop: asyncio.AbstractEventLoop,
    monkeypatch: MonkeyPatch,
):
  test_cases = (
    # test_input, expected
    ({"http_method": "HEAD"}, {"status": 204}),
    ({"http_method": "GET"}, {"status": 204}),
    ({"http_method": "POST"}, {"status": 405}),
  )

  # fmt: off
  async with \
      test_utils.ctx_test_config(monkeypatch=monkeypatch), \
      test_utils.ctx_mqtt_broker(loop=loop, monkeypatch=monkeypatch), \
      test_utils.ctx_app(loop=loop) as app \
  :
    # fmt: on
    client = await aiohttp_client(app)

    for (test_input, expected) in test_cases:
      resp = await client.request(test_input["http_method"], "/ping")
      assert resp.status == expected["status"]
```
This also happens with `# yapf: disable` (#2944).
The same error also occurs without `# fmt: off` being indented:

```python
if True:
# fmt: off
    pass

False
```

I get the following AST diff in the log file:

```diff
--- src
+++ dst
@@ -2,21 +2,21 @@
   body=
     If(
       body=
         Pass(
         )  # /Pass
+        Expr(
+          value=
+            Constant(
+              value=
+                False,  # bool
+            )  # /Constant
+        )  # /Expr
       orelse=
       test=
         Constant(
           value=
             True,  # bool
         )  # /Constant
     )  # /If
-    Expr(
-      value=
-        Constant(
-          value=
-            False,  # bool
-        )  # /Constant
-    )  # /Expr
   type_ignores=
 )  # /Module
\ No newline at end of file
```

This seems to be the same invalid formatting that is provided in this issue's first comment.
I _think_ this can benefit a similar fix as https://github.com/psf/black/pull/3148 (for `# fmt: skip`)

I can take a look at this issue too, but I'll wait for that PR to be reviewed/merged in case changes are needed in that PR.

Resolving patch:
diff --git a/src/black/comments.py b/src/black/comments.py
--- a/src/black/comments.py
+++ b/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
-        and type_repr(parent.type) == "suite"
-        and leaf.type == token.NEWLINE
+        parent is not None and parent.type == syms.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:
-    """Determine if children at a given column have formatting switched on."""
+def children_contains_fmt_on(container: LN, *, preview: bool) -> bool:
+    """Determine if children have formatting switched on."""
     for child in container.children:
-        if (
-            isinstance(child, Node)
-            and first_leaf_column(child) == column
-            or isinstance(child, Leaf)
-            and child.column == column
-        ):
-            if is_fmt_on(child, preview=preview):
-                return True
+        leaf = first_leaf_of(child)
+        if leaf is not None and is_fmt_on(leaf, preview=preview):
+            return True
 
     return False
 
diff --git a/src/black/nodes.py b/src/black/nodes.py
--- a/src/black/nodes.py
+++ b/src/black/nodes.py
@@ -502,12 +502,14 @@ def container_of(leaf: Leaf) -> LN:
     return container
 
 
-def first_leaf_column(node: Node) -> Optional[int]:
-    """Returns the column of the first leaf child of a node."""
-    for child in node.children:
-        if isinstance(child, Leaf):
-            return child.column
-    return None
+def first_leaf_of(node: LN) -> Optional[Leaf]:
+    """Returns the first leaf of the node tree."""
+    if isinstance(node, Leaf):
+        return node
+    if node.children:
+        return first_leaf_of(node.children[0])
+    else:
+        return None
 
 
 def is_arith_like(node: LN) -> bool:


Test patch:
diff --git a/tests/data/simple_cases/fmtonoff5.py b/tests/data/simple_cases/fmtonoff5.py
--- a/tests/data/simple_cases/fmtonoff5.py
+++ b/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