From e867306f59bcf302899c485cc634a47cf30aeff0 Mon Sep 17 00:00:00 2001
From: Aleksei Kniazev <aleksei.kniazev@jetbrains.com>
Date: Tue, 2 Jul 2019 19:25:28 +0300
Subject: [PATCH] IDEA-CR-49176: inline function: optimize imports at
 declaration site after removing function definition (PY-36721)

GitOrigin-RevId: 2bf0ffdeff5a2b5cd084b53b5e1ed76d547d07f2
---
 .../python/refactoring/inline/PyInlineFunctionProcessor.kt   | 5 +++--
 .../optimizeImportsAtDeclarationSite/main.after.py           | 3 +++
 .../inlineFunction/optimizeImportsAtDeclarationSite/main.py  | 3 +++
 .../optimizeImportsAtDeclarationSite/src.after.py            | 2 ++
 .../inlineFunction/optimizeImportsAtDeclarationSite/src.py   | 4 ++++
 .../com/jetbrains/python/refactoring/PyInlineFunctionTest.kt | 5 +++++
 6 files changed, 20 insertions(+), 2 deletions(-)
 create mode 100644 python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/main.after.py
 create mode 100644 python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/main.py
 create mode 100644 python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/src.after.py
 create mode 100644 python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/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 8475933bfbfe..e7d7c5e33032 100644
--- a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
+++ b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
@@ -295,16 +295,17 @@ class PyInlineFunctionProcessor(project: Project,
       .forEach { PyClassRefactoringUtil.optimizeImports(it) }
 
     if (myRemoveDeclaration) {
+      val file = myFunction.containingFile
       val stubFunction = PyiUtil.getPythonStub(myFunction)
       if (stubFunction != null && stubFunction.isWritable) {
         stubFunction.delete()
       }
-      val typingOverloads = PyiUtil.getOverloads(myFunction, TypeEvalContext.userInitiated(myProject, myFunction.containingFile))
+      val typingOverloads = PyiUtil.getOverloads(myFunction, TypeEvalContext.userInitiated(myProject, file))
       if (typingOverloads.isNotEmpty()) {
         typingOverloads.forEach { it.delete() }
-        PyClassRefactoringUtil.optimizeImports(myFunction.containingFile)
       }
       myFunction.delete()
+      PyClassRefactoringUtil.optimizeImports(file)
       dunderAll.forEach { it.element?.delete() }
     }
   }
diff --git a/python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/main.after.py b/python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/main.after.py
new file mode 100644
index 000000000000..ce0685c2d533
--- /dev/null
+++ b/python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/main.after.py
@@ -0,0 +1,3 @@
+from datetime import datetime
+
+print(datetime.now())
diff --git a/python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/main.py b/python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/main.py
new file mode 100644
index 000000000000..74486aaa5761
--- /dev/null
+++ b/python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/main.py
@@ -0,0 +1,3 @@
+from src import foo
+
+f<caret>oo()
\ No newline at end of file
diff --git a/python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/src.after.py b/python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/src.after.py
new file mode 100644
index 000000000000..139597f9cb07
--- /dev/null
+++ b/python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/src.after.py
@@ -0,0 +1,2 @@
+
+
diff --git a/python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/src.py b/python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/src.py
new file mode 100644
index 000000000000..8df9a38e19e2
--- /dev/null
+++ b/python/testData/refactoring/inlineFunction/optimizeImportsAtDeclarationSite/src.py
@@ -0,0 +1,4 @@
+from datetime import datetime
+
+def foo():
+    print(datetime.now())
\ 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 b253b7470fa6..0996ff0e1d68 100644
--- a/python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt
+++ b/python/testSrc/com/jetbrains/python/refactoring/PyInlineFunctionTest.kt
@@ -87,6 +87,11 @@ class PyInlineFunctionTest : PyTestCase() {
   fun testImportedLocally() = doTest(inlineThis = false, remove = true)
   //fun testInlineImportedAs() = doTest(inlineThis = false)
   fun testSelfUsageDetection() = doTest(inlineThis = false, remove = true)
+  fun testOptimizeImportsAtDeclarationSite() {
+    doTest(inlineThis = false, remove = true)
+    val testName = getTestName(true)
+    myFixture.checkResultByFile("src.py", "$testName/src.after.py",true)
+  }
   fun testRemoveFunctionWithStub() {
     doTest(inlineThis = false, remove = true)
     val testName = getTestName(true)