From a2821815af5f5a706c673279d6405e286d6e95b8 Mon Sep 17 00:00:00 2001
From: "Yilei \"Dolee\" Yang" <yileiyang@google.com>
Date: Wed, 14 Dec 2022 17:56:14 -0800
Subject: [PATCH] Fix a crash when a colon line is marked between `# fmt: off`
 and `# fmt: on` (#3439)

---
 CHANGES.md                           |  2 ++
 src/black/comments.py                | 12 +++++++++++-
 tests/data/simple_cases/fmtonoff5.py | 22 ++++++++++++++++++++--
 3 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index f6040359623..1a7c320baf8 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -10,6 +10,8 @@
 
 <!-- Changes that affect Black's stable style -->
 
+- Fix a crash when a colon line is marked between `# fmt: off` and `# fmt: on` (#3439)
+
 ### Preview style
 
 <!-- Changes that affect Black's preview style -->
diff --git a/src/black/comments.py b/src/black/comments.py
index dce83abf1bb..e733dccd844 100644
--- a/src/black/comments.py
+++ b/src/black/comments.py
@@ -232,7 +232,7 @@ def generate_ignored_nodes(
 
         # fix for fmt: on in children
         if children_contains_fmt_on(container, preview=preview):
-            for child in container.children:
+            for index, child in enumerate(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
@@ -241,6 +241,16 @@ def generate_ignored_nodes(
                         # The alternative is to fail the formatting.
                         yield child
                     return
+                if (
+                    child.type == token.INDENT
+                    and index < len(container.children) - 1
+                    and children_contains_fmt_on(
+                        container.children[index + 1], preview=preview
+                    )
+                ):
+                    # This means `# fmt: on` is placed right after an indentation
+                    # level, and we shouldn't swallow the previous INDENT token.
+                    return
                 if children_contains_fmt_on(child, preview=preview):
                     return
                 yield child
diff --git a/tests/data/simple_cases/fmtonoff5.py b/tests/data/simple_cases/fmtonoff5.py
index 71b1381ed0d..181151b6bd6 100644
--- a/tests/data/simple_cases/fmtonoff5.py
+++ b/tests/data/simple_cases/fmtonoff5.py
@@ -64,7 +64,7 @@ async def call(param):
         print ( "This will be formatted" )
 
 
-# Regression test for https://github.com/psf/black/issues/2985
+# Regression test for https://github.com/psf/black/issues/2985.
 class Named(t.Protocol):
     # fmt: off
     @property
@@ -75,6 +75,15 @@ def  this_will_be_formatted ( self, **kwargs ) -> Named: ...
     # fmt: on
 
 
+# Regression test for https://github.com/psf/black/issues/3436.
+if x:
+    return x
+# fmt: off
+elif   unformatted:
+# fmt: on
+    will_be_formatted  ()
+
+
 # output
 
 
@@ -144,7 +153,7 @@ async def call(param):
         print("This will be formatted")
 
 
-# Regression test for https://github.com/psf/black/issues/2985
+# Regression test for https://github.com/psf/black/issues/2985.
 class Named(t.Protocol):
     # fmt: off
     @property
@@ -156,3 +165,12 @@ def this_will_be_formatted(self, **kwargs) -> Named:
         ...
 
     # fmt: on
+
+
+# Regression test for https://github.com/psf/black/issues/3436.
+if x:
+    return x
+# fmt: off
+elif   unformatted:
+    # fmt: on
+    will_be_formatted()