Update order for resolved elements (PY-30942)

python/src/com/jetbrains/python/codeInsight/typing/PyPEP561.kt

@@ -3,11 +3,18 @@
 
 package com.jetbrains.python.codeInsight.typing
 
+import com.intellij.openapi.module.Module
+import com.intellij.openapi.module.ModuleUtilCore
+import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectFileIndex
+import com.intellij.openapi.util.Key
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiDirectory
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFileSystemItem
 import com.intellij.psi.util.QualifiedName
 import com.jetbrains.python.psi.LanguageLevel
+import com.jetbrains.python.psi.PyFile
 import com.jetbrains.python.psi.PyUtil
 import com.jetbrains.python.psi.resolve.PyQualifiedNameResolveContext
 import com.jetbrains.python.psi.resolve.resolveModuleAt
@@ -15,6 +22,9 @@ import com.jetbrains.python.pyi.PyiFile
 import com.jetbrains.python.sdk.PythonSdkType
 
 private const val STUBS_SUFFIX = "-stubs"
+private val STUB_PACKAGE_KEY = Key<Boolean>("PY_STUB_PACKAGE")
+private val INLINE_PACKAGE_KEY = Key<Boolean>("PY_INLINE_PACKAGE")
+private val PEP_561_KEY = Key<Boolean>("PY_PEP_561_KEY")
 
 /**
  * If [name] argument points to element in stub package,
@@ -51,7 +61,11 @@ fun resolveModuleAtStubPackage(name: QualifiedName,
       // check that resolve is running from lib root
       if (virtualFile != null && virtualFile == ProjectFileIndex.getInstance(context.project).getClassRootForFile(virtualFile)) {
         val nameInStubPackage = sequenceOf("$head$STUBS_SUFFIX") + name.components.asSequence().drop(1)
-        return resolveModuleAt(QualifiedName.fromComponents(nameInStubPackage.toList()), item, context).filter(::pyi)
+        return resolveModuleAt(QualifiedName.fromComponents(nameInStubPackage.toList()), item, context)
+          .asSequence()
+          .filter(::pyi)
+          .onEach { it.putUserData(STUB_PACKAGE_KEY, true) }
+          .toList()
       }
     }
   }
@@ -59,6 +73,14 @@ fun resolveModuleAtStubPackage(name: QualifiedName,
   return emptyList()
 }
 
+/**
+ * Filters resolved elements according to their import priority in sys.path and
+ * [PEP 561](https://www.python.org/dev/peps/pep-0561/#type-checker-module-resolution-order) rules.
+ */
+fun filterTopPriorityResults(resolved: List<PsiElement>, module: Module?): List<PsiElement> =
+  if (resolved.all(::isNamespacePackage)) resolved
+  else listOfNotNull(resolved.maxBy { resolvedElementPriority(it, module) })
+
 private fun contextLanguageLevel(context: PyQualifiedNameResolveContext): LanguageLevel {
   context.foothold?.also { return LanguageLevel.forElement(it) }
   context.footholdFile?.also { return LanguageLevel.forElement(it) }
@@ -66,8 +88,98 @@ private fun contextLanguageLevel(context: PyQualifiedNameResolveContext): Langua
   context.sdk?.also { return PythonSdkType.getLanguageLevelForSdk(it) }
   context.effectiveSdk?.also { return PythonSdkType.getLanguageLevelForSdk(it) }
 
-  val moduleSdk = PythonSdkType.findPythonSdk(context.module) ?: return LanguageLevel.getDefault()
-  return PythonSdkType.getLanguageLevelForSdk(moduleSdk)
+  context.module?.also { return PyUtil.getLanguageLevelForModule(it) }
+
+  return LanguageLevel.getDefault()
 }
 
-private fun pyi(element: PsiElement) = element is PyiFile || PyUtil.turnDirIntoInit(element) is PyiFile
+private fun pyi(element: PsiElement) = element is PyiFile || PyUtil.turnDirIntoInit(element) is PyiFile
+
+private fun isNamespacePackage(element: PsiElement): Boolean {
+  if (element is PsiDirectory) {
+    val level = PyUtil.getLanguageLevelForVirtualFile(element.project, element.virtualFile)
+    if (!level.isPython2) {
+      return PyUtil.turnDirIntoInit(element) == null
+    }
+  }
+  return false
+}
+
+/**
+ * See [https://www.python.org/dev/peps/pep-0561/#type-checker-module-resolution-order].
+ */
+private fun resolvedElementPriority(element: PsiElement, module: Module?) = when {
+  isNamespacePackage(element) -> -2
+  isUserFile(element, module) -> if (pyi(element)) 8 else 7
+  isInStubPackage(element, module) -> 6
+  isInTypeShed(element) -> 2
+  isInInlinePackage(element, module) -> 4
+  else -> 0
+}
+
+private fun isUserFile(element: PsiElement, module: Module?) =
+  module != null &&
+  element is PsiFileSystemItem &&
+  element.virtualFile.let {
+    it != null && ModuleUtilCore.moduleContainsFile(module, it, false)
+  }
+
+/**
+ * See [resolveModuleAtStubPackage].
+ */
+private fun isInStubPackage(element: PsiElement, module: Module?) = element.getUserData(STUB_PACKAGE_KEY) == true && isPEP561Enabled(module)
+
+private fun isInTypeShed(element: PsiElement) =
+  pyi(element) && (element as? PsiFileSystemItem)?.virtualFile.let { it != null && PyTypeShed.isInside(it) }
+
+/**
+ * See [https://www.python.org/dev/peps/pep-0561/#packaging-type-information].
+ * Value is cached in element's user data.
+ */
+private fun isInInlinePackage(element: PsiElement, module: Module?): Boolean {
+  val cached = element.getUserData(INLINE_PACKAGE_KEY)
+  if (cached != null) return cached
+
+  val result = !pyi(element) &&
+               (element is PyFile || PyUtil.turnDirIntoInit(element) is PyFile) &&
+               isPEP561Enabled(module) &&
+               isInInlinePackage((element as PsiFileSystemItem).virtualFile, element.project)
+
+  element.putUserData(INLINE_PACKAGE_KEY, result)
+  return result
+}
+
+/**
+ * See [https://www.python.org/dev/peps/pep-0561/].
+ * Value is cached in module's user data.
+ */
+private fun isPEP561Enabled(module: Module?): Boolean {
+  if (module == null) return false
+
+  val cached = module.getUserData(PEP_561_KEY)
+  if (cached != null) return cached
+
+  val result = PyUtil.getLanguageLevelForModule(module).isAtLeast(LanguageLevel.PYTHON37)
+
+  module.putUserData(PEP_561_KEY, result)
+  return result
+}
+
+/**
+ * See [https://www.python.org/dev/peps/pep-0561/#packaging-type-information]
+ */
+private fun isInInlinePackage(file: VirtualFile?, project: Project): Boolean {
+  if (file == null) return false
+
+  val root = ProjectFileIndex.getInstance(project).getClassRootForFile(file) ?: return false
+  var current = if (file.isDirectory) file else file.parent
+
+  while (current != null && current != root && current.isDirectory) {
+    val pyTyped = current.findChild("py.typed")
+    if (pyTyped != null && !pyTyped.isDirectory) return true
+
+    current = current.parent
+  }
+
+  return false
+}

python/src/com/jetbrains/python/psi/PyUtil.java

@@ -31,6 +31,7 @@ import com.intellij.openapi.progress.Task;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.projectRoots.Sdk;
 import com.intellij.openapi.roots.ModuleRootManager;
+import com.intellij.openapi.roots.impl.FilePropertyPusher;
 import com.intellij.openapi.ui.MessageType;
 import com.intellij.openapi.ui.popup.Balloon;
 import com.intellij.openapi.ui.popup.JBPopupFactory;
@@ -664,6 +665,11 @@ public class PyUtil {
     return guessLanguageLevelWithCaching(project);
   }
 
+  @NotNull
+  public static LanguageLevel getLanguageLevelForModule(@NotNull Module module) {
+    return FilePropertyPusher.EP_NAME.findExtensionOrFail(PythonLanguageLevelPusher.class).getImmediateValue(module);
+  }
+
   public static void invalidateLanguageLevelCache(@NotNull Project project) {
     project.putUserData(PythonLanguageLevelPusher.PYTHON_LANGUAGE_LEVEL, null);
   }

python/src/com/jetbrains/python/psi/impl/PythonLanguageLevelPusher.java

@@ -65,7 +65,6 @@ public class PythonLanguageLevelPusher implements FilePropertyPusher<LanguageLev
   static {
     Map<LanguageLevel, LanguageLevel> compatLevels = Maps.newEnumMap(LanguageLevel.class);
     addCompatiblePair(compatLevels, LanguageLevel.PYTHON26, LanguageLevel.PYTHON27);
-    addCompatiblePair(compatLevels, LanguageLevel.PYTHON31, LanguageLevel.PYTHON32);
     addCompatiblePair(compatLevels, LanguageLevel.PYTHON33, LanguageLevel.PYTHON34);
     COMPATIBLE_LEVELS = Maps.immutableEnumMap(compatLevels);
   }
@@ -157,6 +156,7 @@ public class PythonLanguageLevelPusher implements FilePropertyPusher<LanguageLev
   }
 
   @Override
+  @NotNull
   public LanguageLevel getImmediateValue(@NotNull Module module) {
     if (ApplicationManager.getApplication().isUnitTestMode() && LanguageLevel.FORCE_LANGUAGE_LEVEL != null) {
       return LanguageLevel.FORCE_LANGUAGE_LEVEL;

python/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt

@@ -19,6 +19,7 @@ import com.intellij.psi.PsiFileSystemItem
 import com.intellij.psi.PsiManager
 import com.intellij.psi.util.QualifiedName
 import com.jetbrains.python.codeInsight.typing.PyTypeShed
+import com.jetbrains.python.codeInsight.typing.filterTopPriorityResults
 import com.jetbrains.python.codeInsight.typing.resolveModuleAtStubPackage
 import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil
 import com.jetbrains.python.facet.PythonPathContributingFacet
@@ -91,7 +92,7 @@ private fun resolveQualifiedName(name: QualifiedName,
                              resultsFromRoots(name, context),
                              relativeResultsFromSkeletons(name, context)).flatten().distinct()
   val allResults = pythonResults + foreignResults
-  val results = if (name.componentCount > 0) findFirstResults(pythonResults, context.module) + foreignResults else allResults
+  val results = if (name.componentCount > 0) filterTopPriorityResults(pythonResults, context.module) + foreignResults else allResults
 
   if (mayCache) {
     cache?.put(key, results)
@@ -107,7 +108,7 @@ private fun resolveModuleFromRoots(name: QualifiedName, context: PyQualifiedName
   val head = name.removeTail(name.componentCount - 1)
   val nameNoHead = name.removeHead(1)
   return nameNoHead.components.fold(resultsFromRoots(head, context)) { results, component ->
-    findFirstResults(results, context.module)
+    filterTopPriorityResults(results, context.module)
       .asSequence()
       .filterIsInstance<PsiFileSystemItem>()
       .flatMap { resolveModuleAt(QualifiedName.fromComponents(component), it, context).asSequence() }
@@ -227,40 +228,6 @@ fun relativeResultsForStubsFromRoots(name: QualifiedName, context: PyQualifiedNa
   return resultsFromRoots(absoluteName, context.copyWithRelative(-1).copyWithRoots())
 }
 
-/**
- * Filters the results according to their import priority in sys.path.
- */
-private fun findFirstResults(results: List<PsiElement>, module: Module?) =
-    if (results.all(::isNamespacePackage))
-      results
-    else {
-      val result = results.firstOrNull { !isNamespacePackage(it) }
-      val stubFile = results.firstOrNull { it is PyiFile || PyUtil.turnDirIntoInit(it) is PyiFile }
-
-      val resultVFile = (result as? PsiFileSystemItem)?.virtualFile
-      val stubVFile = (stubFile as? PsiFileSystemItem)?.virtualFile
-
-      if (stubFile == null ||
-          module != null &&
-          stubVFile != null && PyTypeShed.isInside(stubVFile) &&
-          resultVFile != null && ModuleUtilCore.moduleContainsFile(module, resultVFile, false)) {
-        listOfNotNull(result)
-      }
-      else {
-        listOfNotNull(stubFile)
-      }
-    }
-
-private fun isNamespacePackage(element: PsiElement): Boolean {
-  if (element is PsiDirectory) {
-    val level = PyUtil.getLanguageLevelForVirtualFile(element.project, element.virtualFile)
-    if (!level.isPython2) {
-      return PyUtil.turnDirIntoInit(element) == null
-    }
-  }
-  return false
-}
-
 private fun resolveWithRelativeLevel(name: QualifiedName, context : PyQualifiedNameResolveContext): List<PsiElement> {
   val footholdFile = context.footholdFile
   if (context.relativeLevel >= 0 && footholdFile != null && !PyUserSkeletonsUtil.isUnderUserSkeletonsDirectory(footholdFile)) {

python/testData/resolve/InlinePackageInsteadTypeShed/lib/multiprocessing/process.py

@@ -0,0 +1 @@
+def current_process() -> str: ...

python/testData/resolve/InlinePackageInsteadTypeShed/main.py

@@ -0,0 +1,4 @@
+from multiprocessing import process
+
+process.current_process
+#        <ref>

python/testData/resolve/StubPackageInsteadInlinePackage/lib/pkg/foo.py

@@ -0,0 +1,2 @@
+def bar(a: str, b: str):
+    return a + b

python/testData/resolve/StubPackageInsteadInlinePackageFullyQName/lib/pkg/foo.py

@@ -0,0 +1,2 @@
+def bar(a: str, b: str):
+    return a + b

python/testData/resolve/StubPackagePy36/lib/pkg/foo.py

@@ -1,2 +0,0 @@
-def bar(a, b):
-    return a + b

python/testData/resolve/TypeShedInsteadPy.py

@@ -0,0 +1,2 @@
+from datetime import MINYEAR
+#                     <ref>

python/testData/resolve/UserPyInsteadStubPackage/lib/pkg/foo.py

@@ -0,0 +1,2 @@
+def bar(a: str, b: str):
+    return a + b

python/testData/resolve/UserPyiInsteadUserPy/foo.pyi

@@ -0,0 +1 @@
+def bar(a: str, b: str) -> str: ...

python/testData/resolve/UserPyiInsteadUserPy/main.py

@@ -0,0 +1,2 @@
+from foo import bar
+#                <ref>

python/testSrc/com/jetbrains/python/Py3ResolveTest.java

@@ -15,8 +15,10 @@
  */
 package com.jetbrains.python;
 
+import com.intellij.openapi.vfs.StandardFileSystems;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiReference;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.testFramework.LightProjectDescriptor;
@@ -657,9 +659,22 @@ public class Py3ResolveTest extends PyResolveTestCase {
   }
 
   // PY-30942
-  public void testStubPackage() {
+  public void testUserPyiInsteadUserPy() {
     myFixture.copyDirectoryToProject("resolve/" + getTestName(false), "");
-    final VirtualFile libDir = myFixture.findFileInTempDir("lib");
+    myFixture.configureByFile("main.py");
+
+    final PsiElement element = PyResolveTestCase.findReferenceByMarker(myFixture.getFile()).resolve();
+    assertInstanceOf(element, PyFunction.class);
+    assertEquals("foo.pyi", element.getContainingFile().getName());
+  }
+
+  // PY-30942
+  public void testUserPyInsteadStubPackage() {
+    final String path = "resolve/" + getTestName(false);
+    myFixture.copyDirectoryToProject(path + "/pkg", "pkg");
+    myFixture.configureByFile(path + "/main.py");
+
+    final VirtualFile libDir = StandardFileSystems.local().findFileByPath(getTestDataPath() + "/" + path + "/lib");
     assertNotNull(libDir);
 
     runWithLanguageLevel(
@@ -668,20 +683,23 @@ public class Py3ResolveTest extends PyResolveTestCase {
         runWithAdditionalClassEntryInSdkRoots(
           libDir,
           () -> {
-            myFixture.configureByFile("main.py");
-
             final PsiElement element = PyResolveTestCase.findReferenceByMarker(myFixture.getFile()).resolve();
             assertInstanceOf(element, PyFunction.class);
-            assertEquals("foo.pyi", element.getContainingFile().getName());
+
+            final PsiFile file = element.getContainingFile();
+            assertEquals("foo.py", file.getName());
+            assertEquals("src", file.getParent().getParent().getName());
           }
         )
     );
   }
 
   // PY-30942
-  public void testStubPackageFullyQName() {
-    myFixture.copyDirectoryToProject("resolve/" + getTestName(false), "");
-    final VirtualFile libDir = myFixture.findFileInTempDir("lib");
+  public void testStubPackageInsteadInlinePackage() {
+    final String path = "resolve/" + getTestName(false);
+    myFixture.configureByFile(path + "/main.py");
+
+    final VirtualFile libDir = StandardFileSystems.local().findFileByPath(getTestDataPath() + "/" + path + "/lib");
     assertNotNull(libDir);
 
     runWithLanguageLevel(
@@ -690,35 +708,65 @@ public class Py3ResolveTest extends PyResolveTestCase {
         runWithAdditionalClassEntryInSdkRoots(
           libDir,
           () -> {
-            myFixture.configureByFile("main.py");
-
             final PsiElement element = PyResolveTestCase.findReferenceByMarker(myFixture.getFile()).resolve();
-            assertInstanceOf(element, PyFunction.class);
-            assertEquals("foo.pyi", element.getContainingFile().getName());
+
+            final PsiFile file = element.getContainingFile();
+            assertEquals("foo.pyi", file.getName());
+            assertEquals("pkg-stubs", file.getParent().getName());
           }
         )
     );
   }
 
   // PY-30942
-  public void testStubPackagePy36() {
-    myFixture.copyDirectoryToProject("resolve/StubPackage", "");
-    final VirtualFile libDir = myFixture.findFileInTempDir("lib");
+  public void testStubPackageInsteadInlinePackageFullyQName() {
+    final String path = "resolve/" + getTestName(false);
+    myFixture.configureByFile(path + "/main.py");
+
+    final VirtualFile libDir = StandardFileSystems.local().findFileByPath(getTestDataPath() + "/" + path + "/lib");
     assertNotNull(libDir);
 
     runWithLanguageLevel(
-      LanguageLevel.PYTHON36,
+      LanguageLevel.PYTHON37,
       () ->
         runWithAdditionalClassEntryInSdkRoots(
           libDir,
           () -> {
-            myFixture.configureByFile("main.py");
-
             final PsiElement element = PyResolveTestCase.findReferenceByMarker(myFixture.getFile()).resolve();
             assertInstanceOf(element, PyFunction.class);
-            assertEquals("foo.py", element.getContainingFile().getName());
+
+            final PsiFile file = element.getContainingFile();
+            assertEquals("foo.pyi", file.getName());
+            assertEquals("pkg-stubs", file.getParent().getName());
           }
         )
     );
   }
+
+  // PY-30942
+  public void testInlinePackageInsteadTypeShed() {
+    final String path = "resolve/" + getTestName(false);
+    myFixture.configureByFile(path + "/main.py");
+
+    final VirtualFile libDir = StandardFileSystems.local().findFileByPath(getTestDataPath() + "/" + path + "/lib");
+    assertNotNull(libDir);
+
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON37,
+      () ->
+        runWithAdditionalClassEntryInSdkRoots(
+          libDir,
+          () -> {
+            final PsiElement element = PyResolveTestCase.findReferenceByMarker(myFixture.getFile()).resolve();
+            assertInstanceOf(element, PyFunction.class);
+            assertEquals("process.py", element.getContainingFile().getName());
+          }
+        )
+    );
+  }
+
+  // PY-30942
+  public void testTypeShedInsteadPy() {
+    assertResolvesTo(PyTargetExpression.class, "MINYEAR", "datetime.pyi");
+  }
 }