Add support to style function definitions with newlines before function stubs (#4318)

* Add support to style function definitions containing newlines before function stubs

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

---------

Co-authored-by: Shantanu <[email protected]>

Author: peterkra25

CHANGES.md

@@ -28,6 +28,9 @@
 
 <!-- Changes to the parser or to version autodetection -->
 
+- Add support to style function definitions containing newlines before function stubs
+  (#4318)
+
 ### Performance
 
 <!-- Changes that improve Black's performance. -->

src/black/linegen.py

@@ -309,8 +309,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
+
+        ...