Support custom number of blank lines between top-level imports and variable definitions (#854)

* Support custom number of blank lines after top-level imports

* add test for non-default value

* change wording of the knob

CHANGELOG

@@ -2,6 +2,12 @@
 # All notable changes to this project will be documented in this file.
 # This project adheres to [Semantic Versioning](http://semver.org/).
 
+## [0.31.0] UNRELEASED
+### Added
+- Add 'BLANK_LINES_BETWEEN_TOP_LEVEL_IMPORTS_AND_VARIABLES' to support setting
+  a custom number of blank lines between top-level imports and variable
+  definitions.
+
 ## [0.30.0] 2020-04-23
 ### Added
 - Added `SPACES_AROUND_LIST_DELIMITERS`, `SPACES_AROUND_DICT_DELIMITERS`,

CONTRIBUTORS

@@ -14,3 +14,4 @@ Eli Bendersky <eliben@google.com>
 Sam Clegg <sbc@google.com>
 Łukasz Langa <ambv@fb.com>
 Oleg Butuzov <butuzov@made.ua>
+Mauricio Herrera Cuadra <mauricio@arareko.net>

README.rst

@@ -401,6 +401,10 @@ Knobs
         class Bar:
             pass
 
+``BLANK_LINES_BETWEEN_TOP_LEVEL_IMPORTS_AND_VARIABLES``
+    Sets the number of desired blank lines between top-level imports and
+    variable definitions. Useful for compatibility with tools like isort.
+
 ``COALESCE_BRACKETS``
     Do not split consecutive brackets. Only relevant when
     ``DEDENT_CLOSING_BRACKETS`` or ``INDENT_CLOSING_BRACKETS``
@@ -567,7 +571,7 @@ Knobs
         [1, 2]
 
     will be formatted as:
-    
+
     .. code-block:: python
 
         [ 1, 2 ]
@@ -665,16 +669,16 @@ Knobs
     ``b`` in this code:
 
     .. code-block:: python
-      
+
       abcdef(
           aReallyLongThing: int,
           b: [Int,
               Int])
-   
+
     With the new knob this is split as:
 
     .. code-block:: python
-      
+
       abcdef(
           aReallyLongThing: int,
           b: [Int, Int])

yapf/yapflib/reformatter.py

@@ -650,6 +650,14 @@ def _CalculateNumberOfNewlines(first_token, indent_depth, prev_uwline,
     # The docstring shouldn't have a newline before it.
     return NO_BLANK_LINES
 
+  if first_token.is_name and not indent_depth:
+    if (prev_uwline.first.value == 'from' or
+        prev_uwline.first.value == 'import'):
+      # Support custom number of blank lines between top-level imports and
+      # variable definitions.
+      return 1 + style.Get(
+          'BLANK_LINES_BETWEEN_TOP_LEVEL_IMPORTS_AND_VARIABLES')
+
   prev_last_token = prev_uwline.last
   if prev_last_token.is_docstring:
     if (not indent_depth and first_token.value in {'class', 'def', 'async'}):

yapf/yapflib/style.py

@@ -104,6 +104,9 @@ _STYLE_HELP = dict(
     BLANK_LINES_AROUND_TOP_LEVEL_DEFINITION=textwrap.dedent("""\
       Number of blank lines surrounding top-level function and class
       definitions."""),
+    BLANK_LINES_BETWEEN_TOP_LEVEL_IMPORTS_AND_VARIABLES=textwrap.dedent("""\
+      Number of blank lines between top-level imports and variable
+      definitions."""),
     COALESCE_BRACKETS=textwrap.dedent("""\
       Do not split consecutive brackets. Only relevant when
       dedent_closing_brackets is set. For example:
@@ -419,6 +422,7 @@ def CreatePEP8Style():
       BLANK_LINE_BEFORE_CLASS_DOCSTRING=False,
       BLANK_LINE_BEFORE_MODULE_DOCSTRING=False,
       BLANK_LINES_AROUND_TOP_LEVEL_DEFINITION=2,
+      BLANK_LINES_BETWEEN_TOP_LEVEL_IMPORTS_AND_VARIABLES=1,
       COALESCE_BRACKETS=False,
       COLUMN_LIMIT=79,
       CONTINUATION_ALIGN_STYLE='SPACE',
@@ -606,6 +610,7 @@ _STYLE_OPTION_VALUE_CONVERTER = dict(
     BLANK_LINE_BEFORE_CLASS_DOCSTRING=_BoolConverter,
     BLANK_LINE_BEFORE_MODULE_DOCSTRING=_BoolConverter,
     BLANK_LINES_AROUND_TOP_LEVEL_DEFINITION=int,
+    BLANK_LINES_BETWEEN_TOP_LEVEL_IMPORTS_AND_VARIABLES=int,
     COALESCE_BRACKETS=_BoolConverter,
     COLUMN_LIMIT=int,
     CONTINUATION_ALIGN_STYLE=_ContinuationAlignStyleStringConverter,

yapftests/reformatter_basic_test.py

@@ -253,6 +253,95 @@ class BasicReformatterTest(yapf_test_helper.YAPFTest):
     uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
     self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
 
+  def testBlankLinesAfterTopLevelImports(self):
+    unformatted_code = textwrap.dedent("""\
+        import foo as bar
+        VAR = 'baz'
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        import foo as bar
+
+        VAR = 'baz'
+        """)
+    uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
+    unformatted_code = textwrap.dedent("""\
+        import foo as bar
+
+        VAR = 'baz'
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        import foo as bar
+
+
+        VAR = 'baz'
+        """)
+    try:
+      style.SetGlobalStyle(
+          style.CreateStyleFromConfig(
+              '{based_on_style: yapf, blank_lines_between_top_level_imports_and_variables: 2}'
+          ))
+      uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+      self.assertCodeEqual(expected_formatted_code,
+                           reformatter.Reformat(uwlines))
+    finally:
+      style.SetGlobalStyle(style.CreateYapfStyle())
+
+    unformatted_code = textwrap.dedent("""\
+        import foo as bar
+        # Some comment
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        import foo as bar
+        # Some comment
+        """)
+    uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
+    unformatted_code = textwrap.dedent("""\
+        import foo as bar
+        class Baz():
+          pass
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        import foo as bar
+
+
+        class Baz():
+          pass
+        """)
+    uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
+    unformatted_code = textwrap.dedent("""\
+        import foo as bar
+        def foobar():
+          pass
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        import foo as bar
+
+
+        def foobar():
+          pass
+        """)
+    uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
+    unformatted_code = textwrap.dedent("""\
+        def foobar():
+          from foo import Bar
+          Bar.baz()
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        def foobar():
+          from foo import Bar
+          Bar.baz()
+        """)
+    uwlines = yapf_test_helper.ParseAndUnwrap(unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
   def testBlankLinesAtEndOfFile(self):
     unformatted_code = textwrap.dedent("""\
         def foobar(): # foo
@@ -421,8 +510,8 @@ class foo(object):\n  \n  def foobar(self):\n    \n    pass\n  \n  def barfoo(se
 
   def testCommentsWithTrailingSpaces(self):
     unformatted_code = textwrap.dedent("""\
-        # Thing 1    
-        # Thing 2    
+        # Thing 1
+        # Thing 2
         """)
     expected_formatted_code = textwrap.dedent("""\
         # Thing 1