From 9ae2649f392a987ca019df6524caf52061a823ab Mon Sep 17 00:00:00 2001
From: Ekaterina Tuzova <Ekaterina.Tuzova@jetbrains.com>
Date: Wed, 1 Aug 2012 14:40:38 +0400
Subject: [PATCH] fixed PY-7089 Insert type assertion: leads to syntactically
 incorrect code when invoked for one-line function

---
 .../intentions/TypeAssertionIntention.java    | 26 ++++++++++++++++++-
 python/testData/intentions/afterDocStub.py    | 14 +++++-----
 python/testData/intentions/beforeDocStub.py   |  4 +--
 .../com/jetbrains/python/PyIntentionTest.java |  9 ++++++-
 4 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java
index ab65d8bda088..fe8fd6f985da 100644
--- a/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java
@@ -3,6 +3,7 @@ package com.jetbrains.python.codeInsight.intentions;
 import com.intellij.codeInsight.CodeInsightUtilBase;
 import com.intellij.codeInsight.intention.IntentionAction;
 import com.intellij.codeInsight.template.*;
+import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
@@ -89,7 +90,30 @@ public class TypeAssertionIntention implements IntentionAction {
         element = parent.addAfter(assertStatement, parentStatement);
       }
       else {
-        element = parent.addBefore(assertStatement, parentStatement);
+        PyStatementList statementList = PsiTreeUtil.getParentOfType(parentStatement, PyStatementList.class);
+        final Document document = editor.getDocument();
+
+        if (statementList != null) {
+          PsiElement statementListParent = statementList.getParent();
+          if (document.getLineNumber(statementList.getTextOffset()) ==
+              document.getLineNumber(statementListParent.getTextOffset())) {
+            final String substring =
+              TextRange.create(statementListParent.getTextOffset(), statementList.getTextOffset()).substring(document.getText());
+            final PyStatement foo =
+              elementGenerator.createFromText(LanguageLevel.forElement(problemElement), PyStatement.class, substring + "\n\t" +
+                                             text + "\n\t" + statementList.getText());
+
+            statementListParent = statementListParent.replace(foo);
+            statementListParent = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(statementListParent);
+            statementList = PsiTreeUtil.findChildOfType(statementListParent, PyStatementList.class);
+            element = statementList.getStatements()[0];
+          }
+          else
+            element = parent.addBefore(assertStatement, parentStatement);
+        }
+        else {
+          element = parent.addBefore(assertStatement, parentStatement);
+        }
       }
 
       int textOffSet = element.getTextOffset();
diff --git a/python/testData/intentions/afterDocStub.py b/python/testData/intentions/afterDocStub.py
index 3cc752bc60e5..261170a0d7ff 100644
--- a/python/testData/intentions/afterDocStub.py
+++ b/python/testData/intentions/afterDocStub.py
@@ -1,9 +1,9 @@
 def foo(a, b):
-  """
+    """
   
-  :param a: 
-  :param b: 
-  :return:
-  """
-  if True:
-    return
\ No newline at end of file
+    :param a:
+    :param b:
+    :return:
+    """
+    if True:
+        return
\ No newline at end of file
diff --git a/python/testData/intentions/beforeDocStub.py b/python/testData/intentions/beforeDocStub.py
index 601eea863368..092d5a7ff1e2 100644
--- a/python/testData/intentions/beforeDocStub.py
+++ b/python/testData/intentions/beforeDocStub.py
@@ -1,3 +1,3 @@
 def fo<caret>o(a, b):
-  if True:
-    return
\ No newline at end of file
+    if True:
+        return
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PyIntentionTest.java b/python/testSrc/com/jetbrains/python/PyIntentionTest.java
index 6bc39186cd7c..ad9c261640fb 100644
--- a/python/testSrc/com/jetbrains/python/PyIntentionTest.java
+++ b/python/testSrc/com/jetbrains/python/PyIntentionTest.java
@@ -32,6 +32,13 @@ public class PyIntentionTest extends PyTestCase {
     }
   }
 
+  private void doTest(String hint, boolean ignoreWhiteSpaces) {
+    myFixture.configureByFile("intentions/before" + getTestName(false) + ".py");
+    final IntentionAction action = myFixture.findSingleIntention(hint);
+    myFixture.launchAction(action);
+    myFixture.checkResultByFile("intentions/after" + getTestName(false) + ".py", ignoreWhiteSpaces);
+  }
+
   /**
    * Ensures that intention with given hint <i>is not</i> active.
    * @param hint
@@ -244,7 +251,7 @@ public class PyIntentionTest extends PyTestCase {
     PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(myFixture.getProject());
     documentationSettings.setFormat(DocStringFormat.REST);
     try {
-      doTest(PyBundle.message("INTN.doc.string.stub"));
+      doTest(PyBundle.message("INTN.doc.string.stub"), true);
     }
     finally {
       documentationSettings.setFormat(DocStringFormat.PLAIN);