PY-9923 Add two blank lines between top-level declarations and surrounding statements

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

@@ -42,6 +42,7 @@ import java.util.List;
 
 import static com.jetbrains.python.formatter.PyCodeStyleSettings.DICT_ALIGNMENT_ON_COLON;
 import static com.jetbrains.python.formatter.PyCodeStyleSettings.DICT_ALIGNMENT_ON_VALUE;
+import static com.jetbrains.python.formatter.PythonFormattingModelBuilder.STATEMENT_OR_DECLARATION;
 
 /**
  * @author yole
@@ -684,10 +685,12 @@ 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);
@@ -700,19 +703,10 @@ public class PyBlock implements ASTBlock {
         }
       }
 
-      if ((childType1 == PyElementTypes.FUNCTION_DECLARATION || childType1 == PyElementTypes.CLASS_DECLARATION)
-          && myNode.getElementType() instanceof PyFileElementType) {
-
-        if (psi2 instanceof PsiComment) {
-          final PsiElement psi3 = PsiTreeUtil.getNextSiblingOfType(psi2, PyElement.class);
-
-          if (psi3 != null) {
-            final IElementType type3 = psi3.getNode().getElementType();
-
-            if (type3 == PyElementTypes.CLASS_DECLARATION || type3 == PyElementTypes.FUNCTION_DECLARATION) {
-              return getBlankLinesForOption(myContext.getPySettings().BLANK_LINES_AROUND_TOP_LEVEL_CLASSES_FUNCTIONS);
-            }
-          }
+      if ((PyElementTypes.CLASS_OR_FUNCTION.contains(childType1) && hasTypeIgnoringPrecedingComments(psi2, STATEMENT_OR_DECLARATION)) ||
+          STATEMENT_OR_DECLARATION.contains(childType1) && hasTypeIgnoringPrecedingComments(psi2, PyElementTypes.CLASS_OR_FUNCTION)) {
+        if (PyUtil.isTopLevel(psi1)) {
+          return getBlankLinesForOption(myContext.getPySettings().BLANK_LINES_AROUND_TOP_LEVEL_CLASSES_FUNCTIONS);
         }
       }
 
@@ -723,6 +717,17 @@ public class PyBlock implements ASTBlock {
     return myContext.getSpacingBuilder().getSpacing(this, child1, child2);
   }
 
+  private static boolean hasTypeIgnoringPrecedingComments(@NotNull PsiElement element, @NotNull TokenSet types) {
+    if (element instanceof PsiComment) {
+      final PsiElement psi3 = PsiTreeUtil.getNextSiblingOfType(element, PyElement.class);
+      if (psi3 != null) {
+        final IElementType type3 = psi3.getNode().getElementType();
+        return types.contains(type3);
+      }
+    }
+    return types.contains(element.getNode().getElementType());
+  }
+
   private Spacing getBlankLinesForOption(final int option) {
     final int blankLines = option + 1;
     return Spacing
@@ -733,10 +738,8 @@ public class PyBlock implements ASTBlock {
     final PyStatement statement = PsiTreeUtil.getParentOfType(myNode.getPsi(), PyStatement.class);
     if (statement != null) {
       final Collection<PyStatementPart> parts = PsiTreeUtil.collectElementsOfType(statement, PyStatementPart.class);
-      if ((parts.size() == 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON) ||
-          (parts.size() > 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON_MULTI_CLAUSE)) {
-        return true;
-      }
+      return (parts.size() == 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON) ||
+             (parts.size() > 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON_MULTI_CLAUSE);
     }
     return false;
   }

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

@@ -40,6 +40,7 @@ import static com.jetbrains.python.PyTokenTypes.*;
 @SuppressWarnings("UseOfSystemOutOrSystemErr")
 public class PythonFormattingModelBuilder implements FormattingModelBuilderEx, CustomFormattingModelBuilder {
   private static final boolean DUMP_FORMATTING_AST = false;
+  public static final TokenSet STATEMENT_OR_DECLARATION = PythonDialectsTokenSetProvider.INSTANCE.getStatementTokens();
 
   @NotNull
   @Override
@@ -75,15 +76,10 @@ public class PythonFormattingModelBuilder implements FormattingModelBuilderEx, C
   protected SpacingBuilder createSpacingBuilder(CodeStyleSettings settings) {
     final IFileElementType file = LanguageParserDefinitions.INSTANCE.forLanguage(PythonLanguage.getInstance()).getFileNodeType();
     final PyCodeStyleSettings pySettings = settings.getCustomSettings(PyCodeStyleSettings.class);
-    final TokenSet STATEMENT_OR_DECLARATION =
-      TokenSet.orSet(PythonDialectsTokenSetProvider.INSTANCE.getStatementTokens(), CLASS_OR_FUNCTION);
 
     final CommonCodeStyleSettings commonSettings = settings.getCommonSettings(PythonLanguage.getInstance());
     return new SpacingBuilder(commonSettings)
-      .betweenInside(STATEMENT_OR_DECLARATION, CLASS_OR_FUNCTION, file).blankLines(pySettings.BLANK_LINES_AROUND_TOP_LEVEL_CLASSES_FUNCTIONS)
-      .betweenInside(CLASS_OR_FUNCTION, STATEMENT_OR_DECLARATION, file).blankLines(pySettings.BLANK_LINES_AROUND_TOP_LEVEL_CLASSES_FUNCTIONS)
       .between(IMPORT_STATEMENTS, TokenSet.andNot(STATEMENT_OR_DECLARATION, IMPORT_STATEMENTS)).blankLines(commonSettings.BLANK_LINES_AFTER_IMPORTS)
-      .betweenInside(CLASS_OR_FUNCTION, CLASS_OR_FUNCTION, file).blankLines(pySettings.BLANK_LINES_AROUND_TOP_LEVEL_CLASSES_FUNCTIONS)
       .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/src/com/jetbrains/python/refactoring/move/PyMoveModuleMembersProcessor.java

@@ -47,6 +47,7 @@ import java.util.List;
 /**
  * Group found usages by moved elements and move each of these elements using {@link PyMoveSymbolProcessor}.
  *
+ * attribute's
  * @author vlan
  * @author Mikhail Golubev
  *

python/testData/formatter/twoLinesBetweenTopLevelStatementAndDeclarationsWithComment.py

@@ -0,0 +1,5 @@
+V1 = 1
+# comment
+class C:
+    pass
+V2 = 2

python/testData/formatter/twoLinesBetweenTopLevelStatementAndDeclarationsWithComment_after.py

@@ -0,0 +1,9 @@
+V1 = 1
+
+
+# comment
+class C:
+    pass
+
+
+V2 = 2

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

@@ -131,10 +131,16 @@ public class PyFormatterTest extends PyTestCase {
     doTest();
   }
 
+  // PY-9923
   public void testTwoLinesBetweenTopLevelDeclarationsWithComment() { // PY-9923
     doTest();
   }
 
+  // PY-9923
+  public void testTwoLinesBetweenTopLevelStatementAndDeclarationsWithComment() {
+    doTest();
+  }
+
   public void testSpecialSlice() {  // PY-1928
     doTest();
   }