Fix an infinite loop when using #fmt: on/off in the middle of an expression or code block

CHANGES.md

@@ -10,6 +10,9 @@
 
 <!-- Changes that affect Black's stable style -->
 
+- Fix an infinite loop when using `#fmt: on/off` in the middle of an expression or code
+  block (#3158)
+
 ### Preview style
 
 <!-- Changes that affect Black's preview style -->

src/black/comments.py

@@ -241,6 +241,7 @@ def generate_ignored_nodes(
                 if contains_fmt_on_at_column(child, leaf.column, preview=preview):
                     return
                 yield child
+            return
         else:
             yield container
             container = container.next_sibling

tests/data/simple_cases/fmtonoff5.py

@@ -0,0 +1,37 @@
+# Regression test for https://github.com/psf/black/issues/3129.
+setup(
+    entry_points={
+        # fmt: off
+        "console_scripts": [
+            "foo-bar"
+            "=foo.bar.:main",
+        # fmt: on
+        ]
+    },
+)
+
+
+# Regression test for https://github.com/psf/black/issues/3129.
+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 a:
+        return True
+    # yapf: enable
+    elif b:
+        return True
+
+    return False