From 1354be2525e4910b8a0d7c46242eae76963db5d2 Mon Sep 17 00:00:00 2001
From: peterkra25 <92831027+peterkra25@users.noreply.github.com>
Date: Tue, 23 Apr 2024 22:19:56 -0400
Subject: [PATCH] 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 <12621235+hauntsaninja@users.noreply.github.com>
---
 CHANGES.md                                |   3 +
 src/black/linegen.py                      |   7 +-
 tests/data/cases/dummy_implementations.py | 105 ++++++++++++++++++++++
 3 files changed, 113 insertions(+), 2 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index c536bff107e..14543add54c 100644
--- a/CHANGES.md
+++ b/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. -->
diff --git a/src/black/linegen.py b/src/black/linegen.py
index 980785e94fa..9d22a7e7854 100644
--- a/src/black/linegen.py
+++ b/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]:
diff --git a/tests/data/cases/dummy_implementations.py b/tests/data/cases/dummy_implementations.py
index 0a52c081bcc..107bc2c506f 100644
--- a/tests/data/cases/dummy_implementations.py
+++ b/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
+
+        ...