From 0bceb70531e48b30bf94df6f498c19dfb7151919 Mon Sep 17 00:00:00 2001
From: Aleksei Kniazev <aleksei.kniazev@jetbrains.com>
Date: Tue, 2 Jul 2019 17:00:36 +0300
Subject: [PATCH] IDEA-CR-49176: select distinct files before optimizing
 imports (PY-36484)

GitOrigin-RevId: 48b187e19ee36fd274c1670ecc388c661f11f774
---
 .../refactoring/inline/PyInlineFunctionProcessor.kt      | 6 +++++-
 .../inlineFunction/importedLocally/main.after.py         | 9 +++++++++
 .../refactoring/inlineFunction/importedLocally/main.py   | 9 +++++++++
 .../refactoring/inlineFunction/importedLocally/src.py    | 4 ++++
 .../inlineFunction/invocationOnImport/main.after.py      | 3 +++
 .../inlineFunction/invocationOnImport/main.py            | 3 +++
 .../refactoring/inlineFunction/invocationOnImport/src.py | 4 ++++
 .../jetbrains/python/refactoring/PyInlineFunctionTest.kt | 2 ++
 8 files changed, 39 insertions(+), 1 deletion(-)
 create mode 100644 python/testData/refactoring/inlineFunction/importedLocally/main.after.py
 create mode 100644 python/testData/refactoring/inlineFunction/importedLocally/main.py
 create mode 100644 python/testData/refactoring/inlineFunction/importedLocally/src.py
 create mode 100644 python/testData/refactoring/inlineFunction/invocationOnImport/main.after.py
 create mode 100644 python/testData/refactoring/inlineFunction/invocationOnImport/main.py
 create mode 100644 python/testData/refactoring/inlineFunction/invocationOnImport/src.py

diff --git a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
index ae5feaa41cbb..77f29fdcea76 100644
--- a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
+++ b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
@@ -282,7 +282,11 @@ class PyInlineFunctionProcessor(project: Project,
       }
     }
 
-    imports.forEach { PyClassRefactoringUtil.optimizeImports(it.element!!.containingFile!!) }
+    imports.asSequence()
+      .map { it.element?.containingFile }
+      .filterNotNull()
+      .distinct()
+      .forEach { PyClassRefactoringUtil.optimizeImports(it) }
 
     if (myRemoveDeclaration) {
       val stubFunction = PyiUtil.getPythonStub(myFunction)
diff --git a/python/testData/refactoring/inlineFunction/importedLocally/main.after.py b/python/testData/refactoring/inlineFunction/importedLocally/main.after.py
new file mode 100644
index 000000000000..a5c72056d0fd
--- /dev/null
+++ b/python/testData/refactoring/inlineFunction/importedLocally/main.after.py
@@ -0,0 +1,9 @@
+z = 1 + 2
+print(z)
+res = z
+
+
+def bar():
+    z = 1 + 2
+    print(z)
+    res1 = z
\ No newline at end of file
diff --git a/python/testData/refactoring/inlineFunction/importedLocally/main.py b/python/testData/refactoring/inlineFunction/importedLocally/main.py
new file mode 100644
index 000000000000..fea302aa2a03
--- /dev/null
+++ b/python/testData/refactoring/inlineFunction/importedLocally/main.py
@@ -0,0 +1,9 @@
+from src import foo
+
+
+res = foo(1, 2)
+
+
+def bar():
+    from src import foo
+    res1 = fo<caret>o(1, 2)
\ No newline at end of file
diff --git a/python/testData/refactoring/inlineFunction/importedLocally/src.py b/python/testData/refactoring/inlineFunction/importedLocally/src.py
new file mode 100644
index 000000000000..3a42708b9a41
--- /dev/null
+++ b/python/testData/refactoring/inlineFunction/importedLocally/src.py
@@ -0,0 +1,4 @@
+def foo(x, y):
+    z = x + y
+    print(z)
+    return z
\ No newline at end of file
diff --git a/python/testData/refactoring/inlineFunction/invocationOnImport/main.after.py b/python/testData/refactoring/inlineFunction/invocationOnImport/main.after.py
new file mode 100644
index 000000000000..0b02931b4ddb
--- /dev/null
+++ b/python/testData/refactoring/inlineFunction/invocationOnImport/main.after.py
@@ -0,0 +1,3 @@
+z = 1 + 2
+print(z)
+res = z
\ No newline at end of file
diff --git a/python/testData/refactoring/inlineFunction/invocationOnImport/main.py b/python/testData/refactoring/inlineFunction/invocationOnImport/main.py
new file mode 100644
index 000000000000..46b5e0bc41b1
--- /dev/null
+++ b/python/testData/refactoring/inlineFunction/invocationOnImport/main.py
@@ -0,0 +1,3 @@
+from src import fo<caret>o
+
+res = foo(1, 2)
\ No newline at end of file
diff --git a/python/testData/refactoring/inlineFunction/invocationOnImport/src.py b/python/testData/refactoring/inlineFunction/invocationOnImport/src.py
new file mode 100644
index 000000000000..3a42708b9a41
--- /dev/null
+++ b/python/testData/refactoring/inlineFunction/invocationOnImport/src.py
@@ -0,0 +1,4 @@
+def foo(x, y):
+    z = x + y
+    print(z)
+    return z
\ 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 561e0f912fbc..f155e08d9a73 100644
--- a/python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt
+++ b/python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt
@@ -83,6 +83,8 @@ class PyInlineFunctionTest : PyTestCase() {
   fun testRemovingDocstringAndTypeComment() = doTest()
   fun testRemovingDocstringOfEmptyFunction() = doTest()
   fun testKeepingComments() = doTest()
+  fun testInvocationOnImport() = doTest(inlineThis = false, remove = true)
+  fun testImportedLocally() = doTest(inlineThis = false, remove = true)
   //fun testInlineImportedAs() = doTest(inlineThis = false)
   fun testRemoveFunctionWithStub() {
     doTest(inlineThis = false, remove = true)