From 3749eba8e1c01ed776d13ec9c5a594f90f9bbc91 Mon Sep 17 00:00:00 2001
From: Aleksei Kniazev <aleksei.kniazev@jetbrains.com>
Date: Mon, 24 Jun 2019 11:55:21 +0300
Subject: [PATCH] IDEA-CR-49176: inline function is made possible for functions
 with lambdas (PY-36373)

GitOrigin-RevId: 4f2603447f760d959d31d5033944bae85ad68bf8
---
 .../refactoring/inline/PyInlineFunctionHandler.kt    |  4 +++-
 .../refactoring/inline/PyInlineFunctionProcessor.kt  |  9 ++++++---
 .../inlineFunction/functionWithLambda/main.after.py  | 12 ++++++++++++
 .../inlineFunction/functionWithLambda/main.py        |  9 +++++++++
 .../python/refactoring/PyInlineFunctionTest.kt       |  1 +
 5 files changed, 31 insertions(+), 4 deletions(-)
 create mode 100644 python/testData/refactoring/inlineFunction/functionWithLambda/main.after.py
 create mode 100644 python/testData/refactoring/inlineFunction/functionWithLambda/main.py

diff --git a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionHandler.kt b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionHandler.kt
index 57ae9d77fa02..e9851107eec1 100644
--- a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionHandler.kt
+++ b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionHandler.kt
@@ -43,7 +43,7 @@ class PyInlineFunctionHandler : InlineActionHandler() {
       isOverride(element, project) -> "refactoring.inline.function.overridden"
       functionScope.hasGlobals() -> "refactoring.inline.function.global"
       functionScope.hasNonLocals() -> "refactoring.inline.function.nonlocal"
-      functionScope.hasNestedScopes() -> "refactoring.inline.function.nested"
+      hasNestedFunction(element) -> "refactoring.inline.function.nested"
       hasNonExhaustiveIfs(element) -> "refactoring.inline.function.interrupts.flow"
       else -> null
     }
@@ -56,6 +56,8 @@ class PyInlineFunctionHandler : InlineActionHandler() {
     }
   }
 
+  private fun hasNestedFunction(function: PyFunction): Boolean = SyntaxTraverser.psiTraverser(function.statementList).traverse().any { it is PyFunction }
+
   private fun hasNonExhaustiveIfs(function: PyFunction): Boolean {
     val returns = mutableListOf<PyReturnStatement>()
 
diff --git a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
index 969bdc344f87..60c23bb2aaaf 100644
--- a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
+++ b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
@@ -181,9 +181,12 @@ class PyInlineFunctionProcessor(project: Project,
       replacementFunction.accept(object : PyRecursiveElementVisitor() {
         override fun visitPyReferenceExpression(node: PyReferenceExpression) {
           if (!node.isQualified) {
-            when (val name = node.name) {
-              in mappedArguments -> argumentReplacements[node] = mappedArguments[name]!!
-              in nameClashes -> nameClashRefs.putValue(name!!, node)
+            val parentLambda = PsiTreeUtil.getParentOfType(node, PyLambdaExpression::class.java)
+            if (parentLambda == null  || parentLambda.parameterList.parameters.none { it.name == node.name }) {
+              when (val name = node.name) {
+                in mappedArguments -> argumentReplacements[node] = mappedArguments[name]!!
+                in nameClashes -> nameClashRefs.putValue(name!!, node)
+              }
             }
           }
           super.visitPyReferenceExpression(node)
diff --git a/python/testData/refactoring/inlineFunction/functionWithLambda/main.after.py b/python/testData/refactoring/inlineFunction/functionWithLambda/main.after.py
new file mode 100644
index 000000000000..483a79bd8ec6
--- /dev/null
+++ b/python/testData/refactoring/inlineFunction/functionWithLambda/main.after.py
@@ -0,0 +1,12 @@
+def foo(lst):
+    a = 1
+    res = map(lambda x: x + a + 1, lst)
+    return list(res)
+
+
+a = 2
+
+lst = [1, 2]
+a1 = 1
+res = map(lambda x: x + a1 + 1, lst)
+x = list(res)
\ No newline at end of file
diff --git a/python/testData/refactoring/inlineFunction/functionWithLambda/main.py b/python/testData/refactoring/inlineFunction/functionWithLambda/main.py
new file mode 100644
index 000000000000..4c34c39660ba
--- /dev/null
+++ b/python/testData/refactoring/inlineFunction/functionWithLambda/main.py
@@ -0,0 +1,9 @@
+def foo(lst):
+    a = 1
+    res = map(lambda x: x + a + 1, lst)
+    return list(res)
+
+
+a = 2
+
+x = fo<caret>o([1, 2])
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt b/python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt
index 74801b235c69..753c97274a40 100644
--- a/python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt
+++ b/python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt
@@ -73,6 +73,7 @@ class PyInlineFunctionTest : PyTestCase() {
   fun testCallFromStaticMethod() = doTest()
   fun testCallFromClassMethod() = doTest()
   fun testComplexQualifier() = doTest()
+  fun testFunctionWithLambda() = doTest()
   fun testRemovingDocstring() = doTest()
   fun testKeepFirstComment() = doTest()
   //fun testInlineImportedAs() = doTest(inlineThis = false)