From 6be74faa86d5a5ef4625e9fe8e42d412b5a0bf4b Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Mon, 19 Nov 2018 16:01:52 +0300
Subject: [PATCH] Do not suggest stub package if it exists locally and resolves
 (PY-31740)

---
 .../codeInsight/typing/PyStubPackages.kt      |  2 +-
 .../typing/PyStubPackagesAdvertiser.kt        | 20 +++++++++++++------
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackages.kt b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackages.kt
index 1e706f78a5c3..66441d32766e 100644
--- a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackages.kt
+++ b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackages.kt
@@ -170,7 +170,7 @@ private fun isUserFile(element: PsiElement, module: Module?) =
 /**
  * See [findStubPackage] and [transferStubPackageMarker].
  */
-private fun isInStubPackage(element: PsiElement) = element.getUserData(STUB_PACKAGE_KEY) == true
+fun isInStubPackage(element: PsiElement) = element.getUserData(STUB_PACKAGE_KEY) == true
 
 private fun isInTypeShed(element: PsiElement) =
   pyi(element) && (element as? PsiFileSystemItem)?.virtualFile.let { it != null && PyTypeShed.isInside(it) }
diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt
index 4e3b204f924f..3b150a347e6d 100644
--- a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt
+++ b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt
@@ -29,6 +29,7 @@ import com.jetbrains.python.packaging.requirement.PyRequirementRelation
 import com.jetbrains.python.psi.PyFile
 import com.jetbrains.python.psi.PyFromImportStatement
 import com.jetbrains.python.psi.PyImportElement
+import com.jetbrains.python.psi.PyReferenceExpression
 import com.jetbrains.python.sdk.PythonSdkType
 import javax.swing.JComponent
 
@@ -69,16 +70,23 @@ class PyStubPackagesAdvertiser : PyInspection() {
 
     override fun visitPyFromImportStatement(node: PyFromImportStatement) {
       super.visitPyFromImportStatement(node)
-      processQName(node.importSourceQName)
+      processImport(node.importSource, node.importSourceQName)
     }
 
     override fun visitPyImportElement(node: PyImportElement) {
       super.visitPyImportElement(node)
-      processQName(node.importedQName)
+      processImport(node.importReferenceExpression, node.importedQName)
     }
 
-    private fun processQName(qualifiedName: QualifiedName?) {
-      qualifiedName?.firstComponent?.let(sources::add)
+    private fun processImport(ref: PyReferenceExpression?, qName: QualifiedName?) {
+      if (qName == null) return
+
+      if (ref != null &&
+          ref.getReference(resolveContext).multiResolve(false).asSequence().mapNotNull { it.element }.any { isInStubPackage(it) }) {
+        return
+      }
+
+      qName.firstComponent?.let(sources::add)
     }
   }
 
@@ -197,9 +205,9 @@ class PyStubPackagesAdvertiser : PyInspection() {
     }
   }
 
-  private fun whiteListedSourcesToProcess(sources: Set<String>): Set<String> = sources.filterTo(mutableSetOf()) { it in WHITE_LIST }
+  private fun whiteListedSourcesToProcess(sources: Set<String>) = sources.filterTo(mutableSetOf()) { it in WHITE_LIST }
 
-  private fun notWhiteListedSourcesToProcess(sources: Set<String>): Set<String> = sources.filterNotTo(mutableSetOf()) { it in WHITE_LIST }
+  private fun notWhiteListedSourcesToProcess(sources: Set<String>) = sources.filterNotTo(mutableSetOf()) { it in WHITE_LIST }
 
   private fun splitIntoNotCachedAndCached(sources: Set<String>,
                                           cache: Cache<String, Set<RepoPackage>>): Pair<Set<String>, Set<RepoPackage>> {