From b0c2bcc9537d238c8a580294ecbc41de465d7f55 Mon Sep 17 00:00:00 2001
From: Richard Si <63936253+ichard26@users.noreply.github.com>
Date: Wed, 1 Dec 2021 18:05:59 -0500
Subject: [PATCH] Treat functions/classes in blocks as if they're nested
 (GH-2472)

* Treat functions/classes in blocks as if they're nested

One curveball is that we still want two preceding newlines before blocks
that are probably logically disconnected. In other words:

    if condition:

        def foo():
            return "hi"
                             # <- aside: this is the goal of this commit
    else:

        def foo():
            return "cya"
                             # <- the two newlines spacing here should stay
                             #    since this probably isn't related
    with open("db.json", encoding="utf-8") as f:
        data = f.read()

Unfortunately that means we have to special case specific clause types
instead of just being able to just for a colon leaf. The hack used here
is to check whether we're adding preceding newlines for a standalone or
dependent clause. "Standalone" being a clause that doesn't need another
clause to be valid (eg. if) and vice versa.

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
---
 CHANGES.md                   |  1 +
 src/black/lines.py           | 24 ++++++++++++--
 src/black_primer/primer.json |  2 +-
 tests/data/function2.py      | 63 ++++++++++++++++++++++++++++++++++++
 4 files changed, 87 insertions(+), 3 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 5b64822..5904291 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -17,6 +17,7 @@
 - Add `flake8-simplify` and `flake8-comprehensions` plugins (#2653)
 - Fix determination of f-string expression spans (#2654)
 - Fix parser error location on invalid syntax in a `match` statement (#2649)
+- Functions and classes in blocks now have more consistent surrounding spacing (#2472)
 
 ## 21.11b1
 
diff --git a/src/black/lines.py b/src/black/lines.py
index 63225c0..f2bdada 100644
--- a/src/black/lines.py
+++ b/src/black/lines.py
@@ -447,11 +447,31 @@ class EmptyLineTracker:
             before = 0
         depth = current_line.depth
         while self.previous_defs and self.previous_defs[-1] >= depth:
-            self.previous_defs.pop()
             if self.is_pyi:
                 before = 0 if depth else 1
             else:
-                before = 1 if depth else 2
+                if depth:
+                    before = 1
+                elif (
+                    not depth
+                    and self.previous_defs[-1]
+                    and current_line.leaves[-1].type == token.COLON
+                    and (
+                        current_line.leaves[0].value
+                        not in ("with", "try", "for", "while", "if", "match")
+                    )
+                ):
+                    # We shouldn't add two newlines between an indented function and
+                    # a dependent non-indented clause. This is to avoid issues with
+                    # conditional function definitions that are technically top-level
+                    # and therefore get two trailing newlines, but look weird and
+                    # inconsistent when they're followed by elif, else, etc. This is
+                    # worse because these functions only get *one* preceding newline
+                    # already.
+                    before = 1
+                else:
+                    before = 2
+            self.previous_defs.pop()
         if current_line.is_decorator or current_line.is_def or current_line.is_class:
             return self._maybe_empty_lines_for_class_or_def(current_line, before)
 
diff --git a/src/black_primer/primer.json b/src/black_primer/primer.json
index 2290d1d..8fe61e8 100644
--- a/src/black_primer/primer.json
+++ b/src/black_primer/primer.json
@@ -116,7 +116,7 @@
     },
     "pyanalyze": {
       "cli_arguments": ["--experimental-string-processing"],
-      "expect_formatting_changes": false,
+      "expect_formatting_changes": true,
       "git_clone_url": "https://github.com/quora/pyanalyze.git",
       "long_checkout": false,
       "py_versions": ["all"]
diff --git a/tests/data/function2.py b/tests/data/function2.py
index cfc259e..5bb36c2 100644
--- a/tests/data/function2.py
+++ b/tests/data/function2.py
@@ -23,6 +23,35 @@ def h():
         pass
     print("Inner defs should breathe a little.")
 
+
+if os.name == "posix":
+    import termios
+    def i_should_be_followed_by_only_one_newline():
+        pass
+elif os.name == "nt":
+    try:
+        import msvcrt
+        def i_should_be_followed_by_only_one_newline():
+            pass
+
+    except ImportError:
+
+        def i_should_be_followed_by_only_one_newline():
+            pass
+
+elif False:
+
+    class IHopeYouAreHavingALovelyDay:
+        def __call__(self):
+            print("i_should_be_followed_by_only_one_newline")
+else:
+
+    def foo():
+        pass
+
+with hmm_but_this_should_get_two_preceding_newlines():
+    pass
+
 # output
 
 def f(
@@ -56,3 +85,37 @@ def h():
         pass
 
     print("Inner defs should breathe a little.")
+
+
+if os.name == "posix":
+    import termios
+
+    def i_should_be_followed_by_only_one_newline():
+        pass
+
+elif os.name == "nt":
+    try:
+        import msvcrt
+
+        def i_should_be_followed_by_only_one_newline():
+            pass
+
+    except ImportError:
+
+        def i_should_be_followed_by_only_one_newline():
+            pass
+
+elif False:
+
+    class IHopeYouAreHavingALovelyDay:
+        def __call__(self):
+            print("i_should_be_followed_by_only_one_newline")
+
+else:
+
+    def foo():
+        pass
+
+
+with hmm_but_this_should_get_two_preceding_newlines():
+    pass
-- 
2.39.5