Relocated implementation for removal of newlines before function stubs with added tests for comments

Author: peterkra25

src/black/linegen.py

@@ -159,8 +159,6 @@ def visit_default(self, node: LN) -> Iterator[Line]:
                 normalize_numeric_literal(node)
             if node.type not in WHITESPACE:
                 self.current_line.append(node)
-            if node.type == token.DOT:
-                node.prefix = node.prefix.lstrip("\n")
         yield from super().visit_default(node)
 
     def visit_test(self, node: Node) -> Iterator[Line]:
@@ -315,8 +313,11 @@ def visit_simple_stmt(self, node: Node) -> Iterator[Line]:
                 yield from self.line(-1)
 
         else:
-            if not node.parent or not is_stub_suite(node.parent):
-                yield from self.line()
+            if node.parent and is_stub_suite(node.parent):
+                node.prefix = ""
+                yield from self.visit_default(node)
+                return
+            yield from self.line()
             yield from self.visit_default(node)
 
     def visit_async_stmt(self, node: Node) -> Iterator[Line]:

tests/data/cases/dummy_implementations.py

@@ -68,6 +68,67 @@ async def async_function(self):
 async def async_function(self):
     ...
 
+class ClassA:
+    def f(self):
+        
+        ...
+
+
+class ClassB:
+    def f(self):
+        
+        
+        
+        
+        
+        
+        
+        
+        
+        ...
+
+
+class ClassC:
+    def f(self):
+        
+        ...
+        # Comment
+
+
+class ClassD:
+    def f(self):# Comment 1
+        
+        ...# Comment 2
+        # Comment 3
+
+
+class ClassE:
+    def f(self):
+        
+        ...
+    def f2(self):
+        print(10)
+
+
+class ClassF:
+    def f(self):
+        
+        ...# Comment 2
+
+
+class ClassG:
+    def f(self):#Comment 1
+        
+        ...# Comment 2
+ 
+       
+class ClassH:
+    def f(self):
+        #Comment
+
+        ...
+
+
 # output
 
 from typing import NoReturn, Protocol, Union, overload
@@ -142,3 +203,47 @@ async def async_function(self): ...
 
 @decorated
 async def async_function(self): ...
+
+
+class ClassA:
+    def f(self): ...
+
+
+class ClassB:
+    def f(self): ...
+
+
+class ClassC:
+    def f(self):
+
+        ...
+        # Comment
+
+
+class ClassD:
+    def f(self):  # Comment 1
+
+        ...  # Comment 2
+        # Comment 3
+
+
+class ClassE:
+    def f(self): ...
+    def f2(self):
+        print(10)
+
+
+class ClassF:
+    def f(self): ...  # Comment 2
+
+
+class ClassG:
+    def f(self):  # Comment 1
+        ...  # Comment 2
+
+
+class ClassH:
+    def f(self):
+        # Comment
+
+        ...