attach trailing comments to statement list (PY-2137)

python/src/com/jetbrains/python/parsing/FollowingCommentBinder.java

@@ -0,0 +1,32 @@
+package com.jetbrains.python.parsing;
+
+import com.intellij.lang.WhitespacesAndCommentsBinder;
+import com.intellij.psi.tree.IElementType;
+import com.jetbrains.python.PyTokenTypes;
+
+import java.util.List;
+
+/**
+* @author yole
+*/
+class FollowingCommentBinder implements WhitespacesAndCommentsBinder {
+  static final FollowingCommentBinder INSTANCE = new FollowingCommentBinder();
+
+  @Override
+  public int getEdgePosition(List<IElementType> tokens, boolean atStreamEdge, TokenTextGetter getter) {
+    int pos = 0;
+    // TODO[yole] handle more cases?
+    while (pos < tokens.size() && tokens.get(pos) == PyTokenTypes.LINE_BREAK) {
+      final CharSequence charSequence = getter.get(pos);
+      if (charSequence.length() == 0 || charSequence.charAt(charSequence.length()-1) != ' ') {
+        break;
+      }
+      pos++;
+      if (pos == tokens.size() || tokens.get(pos) != PyTokenTypes.END_OF_LINE_COMMENT) {
+        break;
+      }
+      pos++;
+    }
+    return pos;
+  }
+}

python/src/com/jetbrains/python/parsing/StatementParsing.java

@@ -720,6 +720,7 @@ public class StatementParsing extends Parsing implements ITokenTypeRemapper {
       }
 
       marker.done(PyElementTypes.STATEMENT_LIST);
+      marker.setCustomEdgeTokenBinders(null, FollowingCommentBinder.INSTANCE);
       if (endMarker != null) {
         endMarker.done(elType);
       }

python/testData/psi/CommentAtEndOfMethod.py

@@ -0,0 +1,3 @@
+def foo():
+    a = 1
+    #comment

python/testData/psi/CommentAtEndOfMethod.txt

@@ -0,0 +1,21 @@
+PyFile:CommentAtEndOfMethod.py
+  PyFunction('foo')
+    PsiElement(Py:DEF_KEYWORD)('def')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:IDENTIFIER)('foo')
+    PyParameterList
+      PsiElement(Py:LPAR)('(')
+      PsiElement(Py:RPAR)(')')
+    PsiElement(Py:COLON)(':')
+    PsiWhiteSpace('\n    ')
+    PyStatementList
+      PyAssignmentStatement
+        PyTargetExpression: a
+          PsiElement(Py:IDENTIFIER)('a')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:EQ)('=')
+        PsiWhiteSpace(' ')
+        PyNumericLiteralExpression
+          PsiElement(Py:INTEGER_LITERAL)('1')
+      PsiWhiteSpace('\n    ')
+      PsiComment(Py:END_OF_LINE_COMMENT)('#comment')

python/testData/psi/CommentBeforeMethod.txt

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

python/testSrc/com/jetbrains/python/PyIndentTest.java

@@ -1,7 +1,6 @@
 package com.jetbrains.python;
 
 import com.intellij.openapi.actionSystem.IdeActions;
-import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.command.CommandProcessor;
 import com.jetbrains.python.fixtures.PyLightFixtureTestCase;
 
@@ -177,6 +176,16 @@ public class PyIndentTest extends PyLightFixtureTestCase {
     );
   }
 
+  public void testIndentAfterTrailingComment() {  // PY-2137
+    doTest("def foo()\n" +
+           "    a = 1\n" +
+           "    #comment<caret>\n",
+           "def foo()\n" +
+           "    a = 1\n" +
+           "    #comment\n" +
+           "    <caret>\n");
+  }
+
   /*
   TODO: formatter core problem?
   public void testAlignListBeforeEquals() throws Exception {

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

@@ -172,6 +172,15 @@ public class PythonLexerTest extends PyLexerTestCase {
            "Py:INDENT", "Py:PASS_KEYWORD");
   }
 
+  public void testDedentAfterComment() { // PY-2137
+    doTest("def foo():\n" +
+           "    pass\n" +
+           "    #comment\n",
+           "Py:DEF_KEYWORD", "Py:SPACE", "Py:IDENTIFIER", "Py:LPAR", "Py:RPAR", "Py:COLON", "Py:STATEMENT_BREAK", "Py:LINE_BREAK",
+           "Py:INDENT", "Py:PASS_KEYWORD", "Py:STATEMENT_BREAK", "Py:LINE_BREAK",
+           "Py:END_OF_LINE_COMMENT", "Py:DEDENT", "Py:LINE_BREAK");
+  }
+
   private static void doTest(String text, String... expectedTokens) {
     doLexerTest(text, new PythonIndentingLexer(), expectedTokens);
   }

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

@@ -227,6 +227,10 @@ public class PythonParsingTest extends ParsingTestCase {
     doTest();
   }
 
+  public void testCommentAtEndOfMethod() { // PY-2137
+    doTest();
+  }
+
   public void doTest() {
     doTest(LanguageLevel.PYTHON25);
   }