From 3ccd8bbd45079c631bcd3c7bf14008af8bb3960e Mon Sep 17 00:00:00 2001
From: peterkra <peterkra@umich.edu>
Date: Sun, 21 Apr 2024 20:37:03 -0400
Subject: [PATCH 1/2] Add support to style function definitions containing
 newlines before function stubs

---
 CHANGES.md                                |  3 +++
 src/black/linegen.py                      |  2 ++
 tests/data/cases/stub_empty_line.py       |  7 +++++++
 tests/data/cases/stub_many_empty_lines.py | 15 +++++++++++++++
 4 files changed, 27 insertions(+)
 create mode 100644 tests/data/cases/stub_empty_line.py
 create mode 100644 tests/data/cases/stub_many_empty_lines.py

diff --git a/CHANGES.md b/CHANGES.md
index af756c79621..2bb270e131e 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -26,6 +26,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 2d9c27a6141..acb4f8d58dc 100644
--- a/src/black/linegen.py
+++ b/src/black/linegen.py
@@ -159,6 +159,8 @@ 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]:
diff --git a/tests/data/cases/stub_empty_line.py b/tests/data/cases/stub_empty_line.py
new file mode 100644
index 00000000000..2d8a7e4f422
--- /dev/null
+++ b/tests/data/cases/stub_empty_line.py
@@ -0,0 +1,7 @@
+class C:
+    def f(self):
+        
+        ...
+# output
+class C:
+    def f(self): ...
diff --git a/tests/data/cases/stub_many_empty_lines.py b/tests/data/cases/stub_many_empty_lines.py
new file mode 100644
index 00000000000..e9521854f3a
--- /dev/null
+++ b/tests/data/cases/stub_many_empty_lines.py
@@ -0,0 +1,15 @@
+class C:
+    def f(self):
+        
+        
+        
+        
+        
+        
+        
+        
+        
+        ...
+# output
+class C:
+    def f(self): ...

From 8dc6217c0a0350ad848d73bcea5ec8cb3dd1c59e Mon Sep 17 00:00:00 2001
From: peterkra <peterkra@umich.edu>
Date: Tue, 23 Apr 2024 21:22:30 -0400
Subject: [PATCH 2/2] Relocated implementation for removal of newlines before
 function stubs with added tests for comments

---
 src/black/linegen.py                      |   9 +-
 tests/data/cases/dummy_implementations.py | 105 ++++++++++++++++++++++
 tests/data/cases/stub_empty_line.py       |   7 --
 tests/data/cases/stub_many_empty_lines.py |  15 ----
 4 files changed, 110 insertions(+), 26 deletions(-)
 delete mode 100644 tests/data/cases/stub_empty_line.py
 delete mode 100644 tests/data/cases/stub_many_empty_lines.py

diff --git a/src/black/linegen.py b/src/black/linegen.py
index acb4f8d58dc..e6b9c3824ce 100644
--- a/src/black/linegen.py
+++ b/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]:
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
+
+        ...
diff --git a/tests/data/cases/stub_empty_line.py b/tests/data/cases/stub_empty_line.py
deleted file mode 100644
index 2d8a7e4f422..00000000000
--- a/tests/data/cases/stub_empty_line.py
+++ /dev/null
@@ -1,7 +0,0 @@
-class C:
-    def f(self):
-        
-        ...
-# output
-class C:
-    def f(self): ...
diff --git a/tests/data/cases/stub_many_empty_lines.py b/tests/data/cases/stub_many_empty_lines.py
deleted file mode 100644
index e9521854f3a..00000000000
--- a/tests/data/cases/stub_many_empty_lines.py
+++ /dev/null
@@ -1,15 +0,0 @@
-class C:
-    def f(self):
-        
-        
-        
-        
-        
-        
-        
-        
-        
-        ...
-# output
-class C:
-    def f(self): ...