From 249c6536c4dff50773f30f222d1f81f0afe41f4c Mon Sep 17 00:00:00 2001
From: "Yilei \"Dolee\" Yang" <yileiyang@google.com>
Date: Tue, 19 Jul 2022 17:57:23 -0700
Subject: [PATCH] Fix an infinite loop when using `# fmt: on/off` ... (#3158)

... in the middle of an expression or code block by adding a missing return.

Co-authored-by: Richard Si <63936253+ichard26@users.noreply.github.com>
Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
---
 CHANGES.md                           |  2 ++
 src/black/comments.py                | 10 +++++++-
 tests/data/simple_cases/fmtonoff5.py | 36 ++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 tests/data/simple_cases/fmtonoff5.py

diff --git a/CHANGES.md b/CHANGES.md
index 90c62de6b98..94c3bdda68e 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -10,6 +10,8 @@
 
 <!-- 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)
 - Fix incorrect handling of `# fmt: skip` on colon `:` lines. (#3148)
 - Comments are no longer deleted when a line had spaces removed around power operators
   (#2874)
diff --git a/src/black/comments.py b/src/black/comments.py
index 522c1a7b88c..7069da16528 100644
--- a/src/black/comments.py
+++ b/src/black/comments.py
@@ -13,7 +13,7 @@
 from blib2to3.pgen2 import token
 
 from black.nodes import first_leaf_column, preceding_leaf, container_of
-from black.nodes import STANDALONE_COMMENT, WHITESPACE
+from black.nodes import CLOSING_BRACKETS, STANDALONE_COMMENT, WHITESPACE
 
 # types
 LN = Union[Leaf, Node]
@@ -227,6 +227,14 @@ def generate_ignored_nodes(
         # fix for fmt: on in children
         if contains_fmt_on_at_column(container, leaf.column, preview=preview):
             for child in container.children:
+                if isinstance(child, Leaf) and is_fmt_on(child, preview=preview):
+                    if child.type in CLOSING_BRACKETS:
+                        # This means `# fmt: on` is placed at a different bracket level
+                        # than `# fmt: off`. This is an invalid use, but as a courtesy,
+                        # we include this closing bracket in the ignored nodes.
+                        # The alternative is to fail the formatting.
+                        yield child
+                    return
                 if contains_fmt_on_at_column(child, leaf.column, preview=preview):
                     return
                 yield child
diff --git a/tests/data/simple_cases/fmtonoff5.py b/tests/data/simple_cases/fmtonoff5.py
new file mode 100644
index 00000000000..746aa41f4e4
--- /dev/null
+++ b/tests/data/simple_cases/fmtonoff5.py
@@ -0,0 +1,36 @@
+# 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