PY-15701 Allow to specify minimum number of blank lines after local imports

Add new option "After local imports" on "Blank lines" tab of Python
code style settings.

python/src/com/jetbrains/python/formatter/PyBlock.java

@@ -685,16 +685,9 @@ public class PyBlock implements ASTBlock {
   public Spacing getSpacing(Block child1, @NotNull Block child2) {
     if (child1 instanceof ASTBlock && child2 instanceof ASTBlock) {
       final ASTNode node1 = ((ASTBlock)child1).getNode();
-      final ASTNode node2 = ((ASTBlock)child2).getNode();
       final PsiElement psi1 = node1.getPsi();
       final PsiElement psi2 = ((ASTBlock)child2).getNode().getPsi();
       final IElementType childType1 = node1.getElementType();
-      final IElementType childType2 = node2.getElementType();
-
-      if (psi1 instanceof PyImportStatementBase && psi2 instanceof PyImportStatementBase &&
-          psi2.getCopyableUserData(IMPORT_GROUP_BEGIN) != null) {
-        return Spacing.createSpacing(0, 0, 2, true, 1);
-      }
 
       final CommonCodeStyleSettings settings = myContext.getSettings();
       if (childType1 == PyTokenTypes.COLON && psi2 instanceof PyStatementList) {
@@ -710,6 +703,21 @@ public class PyBlock implements ASTBlock {
         }
       }
 
+      if (psi1 instanceof PyImportStatementBase) {
+        if (psi2 instanceof PyImportStatementBase &&
+            psi2.getCopyableUserData(IMPORT_GROUP_BEGIN) != null) {
+          return Spacing.createSpacing(0, 0, 2, true, 1);
+        }
+        if (psi2 instanceof PyStatement && !(psi2 instanceof PyImportStatementBase)) {
+          if (PyUtil.isTopLevel(psi1)) {
+            return getBlankLinesForOption(settings.BLANK_LINES_AFTER_IMPORTS);
+          }
+          else {
+            return getBlankLinesForOption(myContext.getPySettings().BLANK_LINES_AFTER_LOCAL_IMPORTS);
+          }
+        }
+      }
+
       if (psi2 instanceof PsiComment && !hasLineBreaksBefore(psi2.getNode(), 1) && myContext.getPySettings().SPACE_BEFORE_NUMBER_SIGN) {
         return Spacing.createSpacing(2, 0, 0, false, 0);
       }
@@ -730,8 +738,9 @@ public class PyBlock implements ASTBlock {
 
   private Spacing getBlankLinesForOption(final int option) {
     final int blankLines = option + 1;
-    return Spacing
-      .createSpacing(0, 0, blankLines, myContext.getSettings().KEEP_LINE_BREAKS, myContext.getSettings().KEEP_BLANK_LINES_IN_DECLARATIONS);
+    return Spacing.createSpacing(0, 0, blankLines,
+                                 myContext.getSettings().KEEP_LINE_BREAKS,
+                                 myContext.getSettings().KEEP_BLANK_LINES_IN_DECLARATIONS);
   }
 
   private boolean needLineBreakInStatement() {

python/src/com/jetbrains/python/formatter/PyCodeStyleSettings.java

@@ -77,6 +77,8 @@ public class PyCodeStyleSettings extends CustomCodeStyleSettings {
   public boolean DICT_NEW_LINE_AFTER_LEFT_BRACE = false;
   public boolean DICT_NEW_LINE_BEFORE_RIGHT_BRACE = false;
 
+  public int BLANK_LINES_AFTER_LOCAL_IMPORTS = 0;
+
   public PyCodeStyleSettings(CodeStyleSettings container) {
     super("Python", container);
   }

python/src/com/jetbrains/python/formatter/PyLanguageCodeStyleSettingsProvider.java

@@ -85,9 +85,17 @@ public class PyLanguageCodeStyleSettingsProvider extends LanguageCodeStyleSettin
                                    "BLANK_LINES_AFTER_IMPORTS",
                                    "KEEP_BLANK_LINES_IN_DECLARATIONS",
                                    "KEEP_BLANK_LINES_IN_CODE");
-      consumer.showCustomOption(PyCodeStyleSettings.class, "BLANK_LINES_AROUND_TOP_LEVEL_CLASSES_FUNCTIONS",
+      consumer.renameStandardOption("BLANK_LINES_AFTER_IMPORTS", "After top-level imports:");
+
+      consumer.showCustomOption(PyCodeStyleSettings.class,
+                                "BLANK_LINES_AROUND_TOP_LEVEL_CLASSES_FUNCTIONS",
                                 "Around top-level classes and functions:",
                                 BLANK_LINES);
+      consumer.showCustomOption(PyCodeStyleSettings.class,
+                                "BLANK_LINES_AFTER_LOCAL_IMPORTS",
+                                "After local imports:",
+                                BLANK_LINES);
+
     }
     else if (settingsType == SettingsType.WRAPPING_AND_BRACES_SETTINGS) {
       consumer.showStandardOptions("RIGHT_MARGIN",

python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java

@@ -79,7 +79,6 @@ public class PythonFormattingModelBuilder implements FormattingModelBuilderEx, C
 
     final CommonCodeStyleSettings commonSettings = settings.getCommonSettings(PythonLanguage.getInstance());
     return new SpacingBuilder(commonSettings)
-      .between(IMPORT_STATEMENTS, TokenSet.andNot(STATEMENT_OR_DECLARATION, IMPORT_STATEMENTS)).blankLines(commonSettings.BLANK_LINES_AFTER_IMPORTS)
       .between(CLASS_DECLARATION, STATEMENT_OR_DECLARATION).blankLines(commonSettings.BLANK_LINES_AROUND_CLASS)
       .between(STATEMENT_OR_DECLARATION, CLASS_DECLARATION).blankLines(commonSettings.BLANK_LINES_AROUND_CLASS)
       .between(FUNCTION_DECLARATION, STATEMENT_OR_DECLARATION).blankLines(commonSettings.BLANK_LINES_AROUND_METHOD)

python/testData/addImport/localFromImport.after.py

@@ -1,6 +1,5 @@
 def func():
     for _ range(10):
         from package.module import foo
-
         foo
 #       <ref>

python/testData/addImport/localImport.after.py

@@ -1,7 +1,6 @@
 def func():
     try:
         import module
-
         module
 #       <ref>
     except:

python/testData/addImport/localImportInlineBranch.after.py

@@ -1,6 +1,5 @@
 def func():
     if True:
         import module
-
         module
 #            <ref>

python/testData/addImport/localImportInlineFunctionBody.after.py

@@ -1,5 +1,4 @@
 def func():
     import module
-
     module
 #           <ref>

python/testData/formatter/noBlankLinesAfterLocalImports.py

@@ -0,0 +1,13 @@
+from pprint import pprint
+VAR = 42
+def foo():
+    import sys
+    import ast, tokenize
+    pass
+
+
+class C:
+    from textwrap import dedent
+    pass
+    import codecs as C
+    pass

python/testData/formatter/noBlankLinesAfterLocalImports_after.py

@@ -0,0 +1,16 @@
+from pprint import pprint
+
+VAR = 42
+
+
+def foo():
+    import sys
+    import ast, tokenize
+    pass
+
+
+class C:
+    from textwrap import dedent
+    pass
+    import codecs as C
+    pass

python/testSrc/com/jetbrains/python/PyFormatterTest.java

@@ -65,6 +65,11 @@ public class PyFormatterTest extends PyTestCase {
     doTest();
   }
 
+  // PY-15701
+  public void testNoBlankLinesAfterLocalImports() {
+    doTest();
+  }
+
   public void testBlankLineBeforeFunction() {
     doTest();
   }