skip preceding comments with same indent also when generating INDENT tokens (PY-2108)

python/src/com/jetbrains/python/lexer/PythonIndentingLexer.java

@@ -286,7 +286,9 @@ public class PythonIndentingLexer extends PythonLexer {
     if (indent > lastIndent) {
       myIndentStack.push(indent);
       myTokenQueue.add(new PendingToken(PyTokenTypes.LINE_BREAK, whiteSpaceStart, whiteSpaceEnd));
-      myTokenQueue.add(new PendingToken(PyTokenTypes.INDENT, whiteSpaceEnd, whiteSpaceEnd));
+      int insertIndex = skipPrecedingCommentsWithIndent(indent, myTokenQueue.size() - 1);
+      int indentOffset = insertIndex == myTokenQueue.size() ? whiteSpaceEnd : myTokenQueue.get(insertIndex).getStart();
+      myTokenQueue.add(insertIndex, new PendingToken(PyTokenTypes.INDENT, indentOffset, indentOffset));
     }
     else if (indent < lastIndent) {
       while (indent < lastIndent) {
@@ -299,20 +301,10 @@ public class PythonIndentingLexer extends PythonLexer {
           insertIndex++;
         }
         else {
-          // insert the DEDENT before previous comments that have the same indent as the current token indent
-          while(insertIndex > 0 && myTokenQueue.get(insertIndex-1) instanceof PendingCommentToken) {
-            final PendingCommentToken commentToken = (PendingCommentToken)myTokenQueue.get(insertIndex - 1);
-            if (commentToken.getIndent() != indent) {
-              break;
-            }
-            insertIndex--;
-            if (insertIndex > 1 &&
-                myTokenQueue.get(insertIndex - 1).getType() == PyTokenTypes.LINE_BREAK &&
-                myTokenQueue.get(insertIndex - 2) instanceof PendingCommentToken) {
-              insertIndex--;
-            }
-            dedentOffset = commentToken.getStart();
-          }
+          insertIndex = skipPrecedingCommentsWithIndent(indent, insertIndex);
+        }
+        if (insertIndex != myTokenQueue.size()) {
+          dedentOffset = myTokenQueue.get(insertIndex).getStart();
         }
         myTokenQueue.add(insertIndex, new PendingToken(PyTokenTypes.DEDENT, dedentOffset, dedentOffset));
       }
@@ -323,6 +315,25 @@ public class PythonIndentingLexer extends PythonLexer {
     }
   }
 
+  private int skipPrecedingCommentsWithIndent(int indent, int index) {
+    // insert the DEDENT before previous comments that have the same indent as the current token indent
+    boolean foundComment = false;
+    while(index > 0 && myTokenQueue.get(index-1) instanceof PendingCommentToken) {
+      final PendingCommentToken commentToken = (PendingCommentToken)myTokenQueue.get(index - 1);
+      if (commentToken.getIndent() != indent) {
+        break;
+      }
+      foundComment = true;
+      index--;
+      if (index > 1 &&
+          myTokenQueue.get(index - 1).getType() == PyTokenTypes.LINE_BREAK &&
+          myTokenQueue.get(index - 2) instanceof PendingCommentToken) {
+        index--;
+      }
+    }
+    return foundComment ? index : myTokenQueue.size();
+  }
+
   private int getNextLineIndent() {
     int indent = 0;
     while (getBaseTokenType() != null && PyTokenTypes.WHITESPACE_OR_LINEBREAK.contains(getBaseTokenType())) {

python/testData/formatter/comment.py

@@ -0,0 +1,6 @@
+def foo(a):
+    if a == 5:
+        # a is 5
+        print 'no'
+
+foo(5)

python/testData/formatter/comment_after.py

@@ -0,0 +1,6 @@
+def foo(a):
+    if a == 5:
+        # a is 5
+        print 'no'
+
+foo(5)

python/testData/psi/CommentAtBeginningOfStatementList.py

@@ -0,0 +1,6 @@
+def foo(a):
+    if a == 5:
+        # a is 5
+        print 'no'
+
+foo(5)

python/testData/psi/CommentAtBeginningOfStatementList.txt

@@ -0,0 +1,45 @@
+PyFile:CommentAtBeginningOfStatementList.py
+  PyFunction('foo')
+    PsiElement(Py:DEF_KEYWORD)('def')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:IDENTIFIER)('foo')
+    PyParameterList
+      PsiElement(Py:LPAR)('(')
+      PyNamedParameter('a')
+        PsiElement(Py:IDENTIFIER)('a')
+      PsiElement(Py:RPAR)(')')
+    PsiElement(Py:COLON)(':')
+    PsiWhiteSpace('\n    ')
+    PyStatementList
+      PyIfStatement
+        PyIfPartIf
+          PsiElement(Py:IF_KEYWORD)('if')
+          PsiWhiteSpace(' ')
+          PyBinaryExpression
+            PyReferenceExpression: a
+              PsiElement(Py:IDENTIFIER)('a')
+            PsiWhiteSpace(' ')
+            PsiElement(Py:EQEQ)('==')
+            PsiWhiteSpace(' ')
+            PyNumericLiteralExpression
+              PsiElement(Py:INTEGER_LITERAL)('5')
+          PsiElement(Py:COLON)(':')
+          PsiWhiteSpace('\n        ')
+          PyStatementList
+            PsiComment(Py:END_OF_LINE_COMMENT)('# a is 5')
+            PsiWhiteSpace('\n        ')
+            PyPrintStatement
+              PsiElement(Py:PRINT_KEYWORD)('print')
+              PsiWhiteSpace(' ')
+              PyStringLiteralExpression: no
+                PsiElement(Py:STRING_LITERAL)(''no'')
+  PsiWhiteSpace('\n\n')
+  PyExpressionStatement
+    PyCallExpression: foo
+      PyReferenceExpression: foo
+        PsiElement(Py:IDENTIFIER)('foo')
+      PyArgumentList
+        PsiElement(Py:LPAR)('(')
+        PyNumericLiteralExpression
+          PsiElement(Py:INTEGER_LITERAL)('5')
+        PsiElement(Py:RPAR)(')')

python/testData/psi/CommentBeforeMethod.txt

@@ -11,7 +11,7 @@ PyFile:CommentBeforeMethod.py
       PyPassStatement
         PsiElement(Py:PASS_KEYWORD)('pass')
       PsiWhiteSpace('\n    ')
-      PsiComment(Py:END_OF_LINE_COMMENT)('# trailing comment\n\n')
+      PsiComment(Py:END_OF_LINE_COMMENT)('# trailing comment')
   PsiWhiteSpace('\n\n')
   PsiComment(Py:END_OF_LINE_COMMENT)('#leading comment')
   PsiWhiteSpace('\n')

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

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

python/testSrc/com/jetbrains/python/PythonLexerTest.java

@@ -105,7 +105,7 @@ public class PythonLexerTest extends PyLexerTestCase {
   public void testIndentedCommentAndCode() {
     doTest("if a:\n #b\n c",
            "Py:IF_KEYWORD", "Py:SPACE", "Py:IDENTIFIER", "Py:COLON", "Py:STATEMENT_BREAK", "Py:LINE_BREAK",
-           "Py:END_OF_LINE_COMMENT", "Py:LINE_BREAK", "Py:INDENT", "Py:IDENTIFIER");
+           "Py:INDENT", "Py:END_OF_LINE_COMMENT", "Py:LINE_BREAK", "Py:IDENTIFIER");
   }
 
   public void testWithNotKeyword() {

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

@@ -231,6 +231,10 @@ public class PythonParsingTest extends ParsingTestCase {
     doTest();
   }
 
+  public void testCommentAtBeginningOfStatementList() {  // PY-2108
+    doTest();
+  }
+
   public void doTest() {
     doTest(LanguageLevel.PYTHON25);
   }