From 892eddacd215d685e136686b7f629ade70adca83 Mon Sep 17 00:00:00 2001
From: otstrel <otstrel@gmail.com>
Date: Fri, 8 May 2020 15:37:17 +0300
Subject: [PATCH] Fix for "# fmt: on" with decorators (#1325)

---
 black.py                | 49 +++++++++++++++++++++++++++++++++--------
 tests/data/fmtonoff4.py | 31 ++++++++++++++++++++++++++
 tests/test_black.py     |  8 +++++++
 3 files changed, 79 insertions(+), 9 deletions(-)
 create mode 100644 tests/data/fmtonoff4.py

diff --git a/black.py b/black.py
index 2df03f7..e55e4fe 100644
--- a/black.py
+++ b/black.py
@@ -3116,18 +3116,49 @@ def generate_ignored_nodes(leaf: Leaf) -> Iterator[LN]:
     """
     container: Optional[LN] = container_of(leaf)
     while container is not None and container.type != token.ENDMARKER:
-        is_fmt_on = False
-        for comment in list_comments(container.prefix, is_endmarker=False):
-            if comment.value in FMT_ON:
-                is_fmt_on = True
-            elif comment.value in FMT_OFF:
-                is_fmt_on = False
-        if is_fmt_on:
+        if fmt_on(container):
             return
 
-        yield container
+        # fix for fmt: on in children
+        if contains_fmt_on_at_column(container, leaf.column):
+            for child in container.children:
+                if contains_fmt_on_at_column(child, leaf.column):
+                    return
+                yield child
+        else:
+            yield container
+            container = container.next_sibling
+
 
-        container = container.next_sibling
+def fmt_on(container: LN) -> bool:
+    is_fmt_on = False
+    for comment in list_comments(container.prefix, is_endmarker=False):
+        if comment.value in FMT_ON:
+            is_fmt_on = True
+        elif comment.value in FMT_OFF:
+            is_fmt_on = False
+    return is_fmt_on
+
+
+def contains_fmt_on_at_column(container: LN, column: int) -> bool:
+    for child in container.children:
+        if (
+            isinstance(child, Node)
+            and first_leaf_column(child) == column
+            or isinstance(child, Leaf)
+            and child.column == column
+        ):
+            if fmt_on(child):
+                return True
+
+    return False
+
+
+def first_leaf_column(node: Node) -> Optional[int]:
+    for child in node.children:
+        if isinstance(child, Leaf):
+            return child.column
+    return None
 
 
 def maybe_make_parens_invisible_in_atom(node: LN, parent: LN) -> bool:
diff --git a/tests/data/fmtonoff4.py b/tests/data/fmtonoff4.py
new file mode 100644
index 0000000..54673c0
--- /dev/null
+++ b/tests/data/fmtonoff4.py
@@ -0,0 +1,31 @@
+# fmt: off
+@test([
+    1, 2,
+    3, 4,
+])
+# fmt: on
+def f(): pass
+
+@test([
+    1, 2,
+    3, 4,
+])
+def f(): pass
+
+# output
+
+# fmt: off
+@test([
+    1, 2,
+    3, 4,
+])
+# fmt: on
+def f():
+    pass
+
+
+@test(
+    [1, 2, 3, 4,]
+)
+def f():
+    pass
diff --git a/tests/test_black.py b/tests/test_black.py
index 7a4a3bb..15b6341 100644
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -632,6 +632,14 @@ class BlackTestCase(unittest.TestCase):
         black.assert_equivalent(source, actual)
         black.assert_stable(source, actual, black.FileMode())
 
+    @patch("black.dump_to_file", dump_to_stderr)
+    def test_fmtonoff4(self) -> None:
+        source, expected = read_data("fmtonoff4")
+        actual = fs(source)
+        self.assertFormatEqual(expected, actual)
+        black.assert_equivalent(source, actual)
+        black.assert_stable(source, actual, black.FileMode())
+
     @patch("black.dump_to_file", dump_to_stderr)
     def test_remove_empty_parentheses_after_class(self) -> None:
         source, expected = read_data("class_blank_parentheses")
-- 
2.39.2