tweak alignment of comments between classes (PY-1598)

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

@@ -6,6 +6,7 @@ import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.*;
 import com.intellij.psi.codeStyle.CodeStyleSettings;
+import com.intellij.psi.impl.source.tree.TreeUtil;
 import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.tree.TokenSet;
 import com.intellij.psi.util.PsiTreeUtil;
@@ -118,7 +119,7 @@ public class PyBlock implements ASTBlock {
     Indent childIndent = Indent.getNoneIndent();
     Alignment childAlignment = null;
     if (childType == PyElementTypes.STATEMENT_LIST || childType == PyElementTypes.IMPORT_ELEMENT) {
-      if (hasLineBreakBefore(child)) {
+      if (hasLineBreaksBefore(child, 1)) {
         childIndent = Indent.getNormalIndent();
       }
     }
@@ -164,18 +165,26 @@ public class PyBlock implements ASTBlock {
         childAlignment = getAlignmentForChildren();
       }
     }
-    try { // maybe enter was pressed and cut us from a previous (nested) statement list
+
+    if (isAfterStatementList(child) && !hasLineBreaksBefore(child, 2)) {  // maybe enter was pressed and cut us from a previous (nested) statement list
+      childIndent = Indent.getNormalIndent();
+    }
+
+    return new PyBlock(child, childAlignment, childIndent, wrap, mySettings);
+  }
+
+  private static boolean isAfterStatementList(ASTNode child) {
+    try {
       PsiElement prev = sure(child.getPsi().getPrevSibling());
       sure(prev instanceof PyStatement);
       PsiElement lastchild = PsiTreeUtil.getDeepestLast(prev);
       sure(lastchild.getParent() instanceof PyStatementList);
-      childIndent = Indent.getNormalIndent();
+      return true;
     }
-    catch (IncorrectOperationException ignored) {
+    catch (IncorrectOperationException e) {
       // not our cup of tea
+      return false;
     }
-
-    return new PyBlock(child, childAlignment, childIndent, wrap, mySettings);
   }
 
   private static boolean needListAlignment(ASTNode child) {
@@ -189,15 +198,22 @@ public class PyBlock implements ASTBlock {
     return true;
   }
 
-  private static boolean hasLineBreakBefore(ASTNode child) {
+  private static boolean hasLineBreaksBefore(ASTNode child, int minCount) {
-    return isWhitespaceWithLineBreaks(child.getTreePrev()) || isWhitespaceWithLineBreaks(child.getFirstChildNode());
+    return isWhitespaceWithLineBreaks(TreeUtil.findLastLeaf(child.getTreePrev()), minCount) ||
+           isWhitespaceWithLineBreaks(child.getFirstChildNode(), minCount);
   }
 
-  private static boolean isWhitespaceWithLineBreaks(ASTNode node) {
+  private static boolean isWhitespaceWithLineBreaks(ASTNode node, int minCount) {
     if (node != null && node.getElementType() == TokenType.WHITE_SPACE) {
       String prevNodeText = node.getText();
-      if (prevNodeText.indexOf('\n') >= 0) {
+      int count = 0;
-        return true;
+      for(int i=0; i<prevNodeText.length(); i++) {
+        if (prevNodeText.charAt(i) == '\n') {
+          count++;
+          if (count == minCount) {
+            return true;
+          }
+        }
       }
     }
     return false;
@@ -412,7 +428,7 @@ public class PyBlock implements ASTBlock {
       // doesn't request childAttributes from the correct block
       while (lastChild != null) {
         IElementType last_type = lastChild.getElementType();
-        if (last_type == PyElementTypes.STATEMENT_LIST && hasLineBreakBefore(lastChild)) {
+        if (last_type == PyElementTypes.STATEMENT_LIST && hasLineBreaksBefore(lastChild, 1)) {
           if (dedentAfterLastStatement((PyStatementList)lastChild.getPsi())) {
             break;
           }

python/testData/formatter/commentBetweenClasses.py

@@ -0,0 +1,9 @@
+class T1(object):
+    def m1(self):
+        pass
+
+# comment about T2
+
+class T2(object):
+    def m2(self):
+        pass

python/testData/formatter/commentBetweenClasses_after.py

@@ -0,0 +1,9 @@
+class T1(object):
+    def m1(self):
+        pass
+
+# comment about T2
+
+class T2(object):
+    def m2(self):
+        pass

python/testData/psi/CommentBetweenClasses.py

@@ -0,0 +1,9 @@
+class T1(object):
+    def m1(self):
+        pass
+
+# comment about T2
+
+class T2(object):
+    def m2(self):
+        pass

python/testData/psi/CommentBetweenClasses.txt

@@ -0,0 +1,56 @@
+PyFile:CommentBetweenClasses.py
+  PyClass: T1
+    PsiElement(Py:CLASS_KEYWORD)('class')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:IDENTIFIER)('T1')
+    PyArgumentList
+      PsiElement(Py:LPAR)('(')
+      PyReferenceExpression: object
+        PsiElement(Py:IDENTIFIER)('object')
+      PsiElement(Py:RPAR)(')')
+    PsiElement(Py:COLON)(':')
+    PsiWhiteSpace('\n    ')
+    PyStatementList
+      PyFunction('m1')
+        PsiElement(Py:DEF_KEYWORD)('def')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:IDENTIFIER)('m1')
+        PyParameterList
+          PsiElement(Py:LPAR)('(')
+          PyNamedParameter('self')
+            PsiElement(Py:IDENTIFIER)('self')
+          PsiElement(Py:RPAR)(')')
+        PsiElement(Py:COLON)(':')
+        PsiWhiteSpace('\n        ')
+        PyStatementList
+          PyPassStatement
+            PsiElement(Py:PASS_KEYWORD)('pass')
+      PsiWhiteSpace('\n\n')
+  PsiComment(Py:END_OF_LINE_COMMENT)('# comment about T2')
+  PsiWhiteSpace('\n\n')
+  PyClass: T2
+    PsiElement(Py:CLASS_KEYWORD)('class')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:IDENTIFIER)('T2')
+    PyArgumentList
+      PsiElement(Py:LPAR)('(')
+      PyReferenceExpression: object
+        PsiElement(Py:IDENTIFIER)('object')
+      PsiElement(Py:RPAR)(')')
+    PsiElement(Py:COLON)(':')
+    PsiWhiteSpace('\n    ')
+    PyStatementList
+      PyFunction('m2')
+        PsiElement(Py:DEF_KEYWORD)('def')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:IDENTIFIER)('m2')
+        PyParameterList
+          PsiElement(Py:LPAR)('(')
+          PyNamedParameter('self')
+            PsiElement(Py:IDENTIFIER)('self')
+          PsiElement(Py:RPAR)(')')
+        PsiElement(Py:COLON)(':')
+        PsiWhiteSpace('\n        ')
+        PyStatementList
+          PyPassStatement
+            PsiElement(Py:PASS_KEYWORD)('pass')

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

@@ -88,6 +88,10 @@ public class PyFormatterTest extends PyLightFixtureTestCase {
     doTest();
   }
 
+  public void testCommentBetweenClasses() { // PY-1598
+    doTest();
+  }
+
   private void doTest() {
     myFixture.configureByFile("formatter/" + getTestName(true) + ".py");
     ApplicationManager.getApplication().runWriteAction(new Runnable() {

python/testSrc/com/jetbrains/python/PythonParsingTest.java

@@ -235,6 +235,10 @@ public class PythonParsingTest extends ParsingTestCase {
     doTest();
   }
 
+  public void testCommentBetweenClasses() {  // PY-1598
+    doTest();
+  }
+
   public void doTest() {
     doTest(LanguageLevel.PYTHON25);
   }