Fix another case where we format dummy implementation for non-functions/classes (#4103)

Author: JelleZijlstra

CHANGES.md

@@ -21,7 +21,7 @@
   docstring (#4060)
 - Fix crash in preview mode when using a short `--line-length` (#4086)
 - Keep suites consisting of only an ellipsis on their own lines if they are not
-  functions or class definitions (#4066)
+  functions or class definitions (#4066) (#4103)
 
 ### Configuration
 

src/black/linegen.py

@@ -42,6 +42,7 @@
     is_atom_with_invisible_parens,
     is_docstring,
     is_empty_tuple,
+    is_function_or_class,
     is_lpar_token,
     is_multiline_string,
     is_name_token,
@@ -299,11 +300,12 @@ def visit_simple_stmt(self, node: Node) -> Iterator[Line]:
                 wrap_in_parentheses(node, child, visible=False)
             prev_type = child.type
 
-        is_suite_like = node.parent and node.parent.type in STATEMENT
-        if is_suite_like:
-            if (
-                self.mode.is_pyi or Preview.dummy_implementations in self.mode
-            ) and is_stub_body(node):
+        if node.parent and node.parent.type in STATEMENT:
+            if Preview.dummy_implementations in self.mode:
+                condition = is_function_or_class(node.parent)
+            else:
+                condition = self.mode.is_pyi
+            if condition and is_stub_body(node):
                 yield from self.visit_default(node)
             else:
                 yield from self.line(+1)

src/black/nodes.py

@@ -736,15 +736,18 @@ def is_funcdef(node: Node) -> bool:
     return node.type == syms.funcdef
 
 
+def is_function_or_class(node: Node) -> bool:
+    return node.type in {syms.funcdef, syms.classdef, syms.async_funcdef}
+
+
 def is_stub_suite(node: Node, mode: Mode) -> bool:
     """Return True if `node` is a suite with a stub body."""
-    if node.parent is not None:
-        if Preview.dummy_implementations in mode and node.parent.type not in (
-            syms.funcdef,
-            syms.async_funcdef,
-            syms.classdef,
-        ):
-            return False
+    if (
+        node.parent is not None
+        and Preview.dummy_implementations in mode
+        and not is_function_or_class(node.parent)
+    ):
+        return False
 
     # If there is a comment, we want to keep it.
     if node.prefix.strip():

tests/data/cases/preview_dummy_implementations.py

@@ -56,6 +56,8 @@ def has_comment():
 if some_condition:
     ...
 
+if already_dummy: ...
+
 # output
 
 from typing import NoReturn, Protocol, Union, overload
@@ -116,3 +118,6 @@ def has_comment(): ...  # still a dummy
 
 if some_condition:
     ...
+
+if already_dummy:
+    ...