From 2f2da7bc032207026e336848b20840270c120afa Mon Sep 17 00:00:00 2001
From: Aleksei Kniazev <aleksei.kniazev@jetbrains.com>
Date: Fri, 28 Jun 2019 00:21:57 +0300
Subject: [PATCH] IDEA-CR-49176: function inline: only inline all is available
 if it is called on import reference (PY-36658)

GitOrigin-RevId: 79d6ac71e03f946140b52b321571bdb04f7550ee
---
 .../inline/PyInlineFunctionDialog.kt          | 26 +++++++++++++++++--
 .../inline/PyInlineFunctionProcessor.kt       | 10 ++++---
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionDialog.kt b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionDialog.kt
index 01b746f77d0c..2dce785cf2f8 100644
--- a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionDialog.kt
+++ b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionDialog.kt
@@ -3,12 +3,17 @@ package com.jetbrains.python.refactoring.inline
 
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiNameIdentifierOwner
 import com.intellij.psi.PsiReference
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.refactoring.inline.InlineOptionsDialog
 import com.jetbrains.python.PyBundle
 import com.jetbrains.python.psi.PyFunction
 import com.jetbrains.python.psi.PyImportStatementBase
+import com.jetbrains.python.psi.PyReferenceExpression
+import com.jetbrains.python.pyi.PyiUtil
 
 /**
  * @author Aleksei.Kniazev
@@ -22,13 +27,17 @@ class PyInlineFunctionDialog(project: Project,
   private val myNumberOfOccurrences: Int = getNumberOfOccurrences(myFunction)
 
   init {
-    myInvokedOnReference = myReference != null
+    myInvokedOnReference = if (myReference != null) {
+      val expression = myReference.element as PyReferenceExpression
+      PsiTreeUtil.getParentOfType(expression, PyImportStatementBase::class.java) == null
+    } else false
     title = if (isMethod) "Inline method $myFunctionName" else "Inline function $myFunctionName"
     init()
   }
 
   override fun doAction() {
-    invokeRefactoring(PyInlineFunctionProcessor(myProject, myEditor, myFunction, myReference, isInlineThisOnly, !isKeepTheDeclaration))
+    val originalFunction = PyiUtil.getOriginalElement(myFunction) as PyFunction?
+    invokeRefactoring(PyInlineFunctionProcessor(myProject, myEditor, originalFunction ?: myFunction, myReference, isInlineThisOnly, !isKeepTheDeclaration))
   }
 
   override fun getNameLabelText(): String {
@@ -50,4 +59,17 @@ class PyInlineFunctionDialog(project: Project,
   override fun ignoreOccurrence(reference: PsiReference): Boolean {
     return PsiTreeUtil.getParentOfType(reference.element, PyImportStatementBase::class.java) == null
   }
+
+  override fun getNumberOfOccurrences(nameIdentifierOwner: PsiNameIdentifierOwner?): Int {
+    // TODO: this override will be redundant after PY-26881 is fixed and should be deleted then
+    val originalNum = super.getNumberOfOccurrences(nameIdentifierOwner)
+    val stubOrImplementation = if (PyiUtil.isInsideStub(myFunction)) PyiUtil.getOriginalElement(myFunction) else PyiUtil.getPythonStub(myFunction)
+    if (originalNum != -1 && stubOrImplementation != null) {
+      val fromDeclaration = ReferencesSearch.search(stubOrImplementation, GlobalSearchScope.projectScope(myProject)).asSequence()
+        .filter(this::ignoreOccurrence)
+        .count()
+      return originalNum + fromDeclaration
+    }
+    return originalNum
+  }
 }
\ No newline at end of file
diff --git a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
index 960658a1ba06..4c59047b0387 100644
--- a/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
+++ b/python/src/com/jetbrains/python/refactoring/inline/PyInlineFunctionProcessor.kt
@@ -95,12 +95,16 @@ class PyInlineFunctionProcessor(project: Project,
   override fun findUsages(): Array<UsageInfo> {
     if (myInlineThis) {
       val element = myReference!!.element as PyReferenceExpression
-      val isImport = PsiTreeUtil.getParentOfType(element, PyImportStatementBase::class.java) != null
       val localImport = PyResolveUtil.resolveLocally(ScopeUtil.getScopeOwner(element)!!, element.name!!).firstOrNull { it is PyImportElement }
-      return if (!isImport && localImport != null) arrayOf(UsageInfo(element), UsageInfo(localImport)) else arrayOf(UsageInfo(element))
+      return if (localImport != null) arrayOf(UsageInfo(element), UsageInfo(localImport)) else arrayOf(UsageInfo(element))
     }
 
-    return ReferencesSearch.search(myFunction, myRefactoringScope).asSequence()
+    // TODO: replace with PyRefactoringUtil#findUsages after PY-26881 and PY-36493 are fixed
+    var references = ReferencesSearch.search(myFunction, myRefactoringScope).findAll().asSequence()
+    PyiUtil.getPythonStub(myFunction)?.let { stub ->
+      references += ReferencesSearch.search(stub, myRefactoringScope).asSequence()
+    }
+    return references
       .distinct()
       .map(PsiReference::getElement)
       .map(::UsageInfo)