From 8683b66a58bb96bc8841fcb8dc4f6add8a6448ac Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Tue, 16 Mar 2010 21:17:57 +0300
Subject: [PATCH] fix indent after inserting newline after comment in function
 body (PY-641)

---
 .../com/jetbrains/python/formatter/PyBlock.java   | 15 +++++++++------
 .../inspections/AddToImportFromList_after.py      |  2 +-
 .../com/jetbrains/python/PyIndentTest.java        |  4 ++++
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/python/src/com/jetbrains/python/formatter/PyBlock.java b/python/src/com/jetbrains/python/formatter/PyBlock.java
index d94d2bfe4c55..37d996e54a3a 100644
--- a/python/src/com/jetbrains/python/formatter/PyBlock.java
+++ b/python/src/com/jetbrains/python/formatter/PyBlock.java
@@ -33,7 +33,7 @@ public class PyBlock implements ASTBlock {
   private final ASTNode _node;
   private final Wrap _wrap;
   private final CodeStyleSettings mySettings;
-  private List<Block> _subBlocks = null;
+  private List<PyBlock> _subBlocks = null;
   private final Alignment _childListAlignment;
   private final TokenSet _listElementTypes;
   private static final boolean DUMP_FORMATTING_BLOCKS = false;
@@ -78,8 +78,8 @@ public class PyBlock implements ASTBlock {
     return new ArrayList<Block>(_subBlocks);
   }
 
-  private List<Block> buildSubBlocks() {
-    List<Block> blocks = new ArrayList<Block>();
+  private List<PyBlock> buildSubBlocks() {
+    List<PyBlock> blocks = new ArrayList<PyBlock>();
     for (ASTNode child = _node.getFirstChildNode(); child != null; child = child.getTreeNext()) {
 
       IElementType childType = child.getElementType();
@@ -263,7 +263,7 @@ public class PyBlock implements ASTBlock {
         return ChildAttributes.DELEGATE_TO_PREV_CHILD;
       }
       
-      PyBlock insertAfterBlock = (PyBlock)_subBlocks.get(newChildIndex - 1);
+      PyBlock insertAfterBlock = _subBlocks.get(newChildIndex - 1);
 
       ASTNode prevNode = insertAfterBlock.getNode();
       PsiElement prevElt = prevNode.getPsi();
@@ -327,10 +327,13 @@ public class PyBlock implements ASTBlock {
   private Indent getChildIndent(int newChildIndex) {
     ASTNode lastChild = getLastNonSpaceChild(_node, false);
     if (lastChild != null && lastChild.getElementType() == PyElementTypes.STATEMENT_LIST && _subBlocks.size() >= newChildIndex) {
-      PyBlock insertAfterBlock = (PyBlock)_subBlocks.get(newChildIndex - 1);
+      int prevIndex = newChildIndex-1;
+      while (prevIndex > 0 && _subBlocks.get(prevIndex).getNode().getElementType() == PyTokenTypes.END_OF_LINE_COMMENT) {
+        prevIndex--;
+      }
+      PyBlock insertAfterBlock = _subBlocks.get(prevIndex);
       ASTNode afterNode = insertAfterBlock.getNode();
 
-
       // handle pressing Enter after colon and before first statement in
       // existing statement list
       if (afterNode.getElementType() == PyElementTypes.STATEMENT_LIST || afterNode.getElementType() == PyTokenTypes.COLON) {
diff --git a/python/testData/inspections/AddToImportFromList_after.py b/python/testData/inspections/AddToImportFromList_after.py
index 68a2fc4de61f..ad6b4b707e87 100644
--- a/python/testData/inspections/AddToImportFromList_after.py
+++ b/python/testData/inspections/AddToImportFromList_after.py
@@ -1,3 +1,3 @@
-from AddToImportFromFoo import bar,foo
+from AddToImportFromFoo import bar, foo
 
 foo # must get imported
diff --git a/python/testSrc/com/jetbrains/python/PyIndentTest.java b/python/testSrc/com/jetbrains/python/PyIndentTest.java
index a5067723ea8a..69d25765d860 100644
--- a/python/testSrc/com/jetbrains/python/PyIndentTest.java
+++ b/python/testSrc/com/jetbrains/python/PyIndentTest.java
@@ -136,6 +136,10 @@ public class PyIndentTest extends PyLightFixtureTestCase {
     doTest("{<caret>}", "{\n    <caret>\n}");
   }
 
+  public void testIndentAfterComment() {   // PY-641
+    doTest("def foo():\n    #some_call()<caret>\n    another_call()", "def foo():\n    #some_call()\n    <caret>\n    another_call()");
+  }
+
   /*
   TODO: formatter core problem?
   public void testAlignListBeforeEquals() throws Exception {