From 959bb0e3e72fdbe6b9344bbfb667b2db04c35f06 Mon Sep 17 00:00:00 2001
From: Richard Si <63936253+ichard26@users.noreply.github.com>
Date: Thu, 30 Dec 2021 19:04:52 -0500
Subject: [PATCH] Stubs: preserve blank line between attributes and methods

---
 CHANGES.md          |  2 ++
 src/black/lines.py  | 21 +++++++++++++++++----
 tests/data/stub.pyi | 27 +++++++++++++++++++++++++--
 3 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index cb637d94c11..8bb96bb1f29 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -17,6 +17,8 @@
 - Tuple unpacking on `return` and `yield` constructs now implies 3.8+ (#2700)
 - Unparenthesized tuples on annotated assignments (e.g
   `values: Tuple[int, ...] = 1, 2, 3`) now implies 3.8+ (#2708)
+- For stubs, one blank line between class attributes and methods is now kept if there's
+  at least one pre-existing blank line (#2736)
 
 ### Packaging
 
diff --git a/src/black/lines.py b/src/black/lines.py
index f2bdada008a..d8617d83bf7 100644
--- a/src/black/lines.py
+++ b/src/black/lines.py
@@ -448,7 +448,14 @@ def _maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
         depth = current_line.depth
         while self.previous_defs and self.previous_defs[-1] >= depth:
             if self.is_pyi:
-                before = 0 if depth else 1
+                assert self.previous_line is not None
+                if depth and not current_line.is_def and self.previous_line.is_def:
+                    # Empty lines between attributes and methods should be preserved.
+                    before = min(1, before)
+                elif depth:
+                    before = 0
+                else:
+                    before = 1
             else:
                 if depth:
                     before = 1
@@ -532,9 +539,15 @@ def _maybe_empty_lines_for_class_or_def(
             elif (
                 current_line.is_def or current_line.is_decorator
             ) and not self.previous_line.is_def:
-                # Blank line between a block of functions (maybe with preceding
-                # decorators) and a block of non-functions
-                newlines = 1
+                if not current_line.depth:
+                    # Blank line between a block of functions (maybe with preceding
+                    # decorators) and a block of non-functions
+                    newlines = 1
+                else:
+                    # In classes empty lines between attributes and methods should
+                    # be preserved. The +1 offset is to negate the -1 done later as
+                    # this function is indented.
+                    newlines = min(2, before + 1)
             else:
                 newlines = 0
         else:
diff --git a/tests/data/stub.pyi b/tests/data/stub.pyi
index 94ba852e018..9a246211284 100644
--- a/tests/data/stub.pyi
+++ b/tests/data/stub.pyi
@@ -2,32 +2,55 @@ X: int
 
 def f(): ...
 
+
+class D: 
+    ...
+
+
 class C:
     ...
 
 class B:
-    ...
+    this_lack_of_newline_should_be_kept: int
+    def b(self) -> None: ...
+
+    but_this_newline_should_also_be_kept: int
 
 class A:
+    attr: int
+    attr2: str
+
     def f(self) -> int:
         ...
 
     def g(self) -> str: ...
 
+
+
 def g():
     ...
 
 def h(): ...
 
+
 # output
 X: int
 
 def f(): ...
 
+class D: ...
 class C: ...
-class B: ...
+
+class B:
+    this_lack_of_newline_should_be_kept: int
+    def b(self) -> None: ...
+
+    but_this_newline_should_also_be_kept: int
 
 class A:
+    attr: int
+    attr2: str
+
     def f(self) -> int: ...
     def g(self) -> str: ...