Fix for "# fmt: on" with decorators (#1325)

[etc/vim.git] / black.py

diff --git a/black.py b/black.py
index 2df03f7bacd0adeee573c8a7d2c9228af5a34b30..e55e4fe197224081bc0206a43a10bc307966f23b 100644 (file)
--- 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: