From a2843202bd4ccf9048bd414768d6c72be10efd6d Mon Sep 17 00:00:00 2001
From: Ekaterina Tuzova <Ekaterina.Tuzova@jetbrains.com>
Date: Tue, 26 Jun 2012 20:44:58 +0400
Subject: [PATCH] fixed PY-6835 Convert lambda to function: leads to
 syntactically incorrect code when lambda expression is the only statement in
 function

---
 .../PyConvertLambdaToFunctionIntention.java   | 25 ++++++-------------
 .../afterConvertLambdaToFunction2.py          |  5 ++++
 .../beforeConvertLambdaToFunction2.py         |  2 ++
 .../com/jetbrains/python/PyIntentionTest.java |  4 +++
 4 files changed, 18 insertions(+), 18 deletions(-)
 create mode 100644 python/testData/intentions/afterConvertLambdaToFunction2.py
 create mode 100644 python/testData/intentions/beforeConvertLambdaToFunction2.py

diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertLambdaToFunctionIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertLambdaToFunctionIntention.java
index 4fa5480b4c84..2e31213513d5 100644
--- a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertLambdaToFunctionIntention.java
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertLambdaToFunctionIntention.java
@@ -70,25 +70,14 @@ public class PyConvertLambdaToFunctionIntention extends BaseIntentionAction {
       functionBuilder.statement("return " + body.getText());
       PyFunction function = functionBuilder.buildFunction(project, LanguageLevel.getDefault());
 
-      PyFunction parentFunction = PsiTreeUtil.getTopmostParentOfType(lambdaExpression, PyFunction.class);
-      if (parentFunction != null ) {
-        PyClass parentClass = PsiTreeUtil.getTopmostParentOfType(parentFunction, PyClass.class);
-        if (parentClass != null) {
-          PsiElement classParent = parentClass.getParent();
-          function = (PyFunction)classParent.addBefore(function, parentClass);
-        } else {
-          PsiElement funcParent = parentFunction.getParent();
-          function = (PyFunction)funcParent.addBefore(function, parentFunction);
-        }
-      } else {
-        PyStatement statement = PsiTreeUtil.getTopmostParentOfType(lambdaExpression,
-                                                                   PyStatement.class);
-        if (statement != null) {
-          PsiElement statementParent = statement.getParent();
-          if (statementParent != null)
-            function = (PyFunction)statementParent.addBefore(function, statement);
-        }
+      final PyStatement statement = PsiTreeUtil.getParentOfType(lambdaExpression,
+                                                                 PyStatement.class);
+      if (statement != null) {
+        final PsiElement statementParent = statement.getParent();
+        if (statementParent != null)
+          function = (PyFunction)statementParent.addBefore(function, statement);
       }
+
       function = CodeInsightUtilBase
         .forcePsiPostprocessAndRestoreElement(function);
 
diff --git a/python/testData/intentions/afterConvertLambdaToFunction2.py b/python/testData/intentions/afterConvertLambdaToFunction2.py
new file mode 100644
index 000000000000..113f460f0a8c
--- /dev/null
+++ b/python/testData/intentions/afterConvertLambdaToFunction2.py
@@ -0,0 +1,5 @@
+def transform():
+  def function(x):
+      return x + 5
+
+  return function
\ No newline at end of file
diff --git a/python/testData/intentions/beforeConvertLambdaToFunction2.py b/python/testData/intentions/beforeConvertLambdaToFunction2.py
new file mode 100644
index 000000000000..cb516e67862d
--- /dev/null
+++ b/python/testData/intentions/beforeConvertLambdaToFunction2.py
@@ -0,0 +1,2 @@
+def transform():
+  return lambda <caret>x: x + 5
\ 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 5c247606c57a..7389a52b9c61 100644
--- a/python/testSrc/com/jetbrains/python/PyIntentionTest.java
+++ b/python/testSrc/com/jetbrains/python/PyIntentionTest.java
@@ -207,6 +207,10 @@ public class PyIntentionTest extends PyTestCase {
     doNegativeTest(PyBundle.message("INTN.convert.lambda.to.function"));
   }
 
+  public void testConvertLambdaToFunction2() {    //PY-6835
+    doTest(PyBundle.message("INTN.convert.lambda.to.function"));
+  }
+
   public void testConvertVariadicParam() { //PY-2264
     doTest(PyBundle.message("INTN.convert.variadic.param"));
   }