Add resolve to stub packages (PY-30942)

2026-09-27 10:03:11 +07:00 · 2018-10-09 20:57:55 +03:00
parent 50c486f260
commit 7863e90505
19 changed files with 181 additions and 10 deletions
@@ -0,0 +1,73 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+@file:JvmName("PyPEP561")
+
+package com.jetbrains.python.codeInsight.typing
+
+import com.intellij.openapi.roots.ProjectFileIndex
+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.PyUtil
+import com.jetbrains.python.psi.resolve.PyQualifiedNameResolveContext
+import com.jetbrains.python.psi.resolve.resolveModuleAt
+import com.jetbrains.python.pyi.PyiFile
+import com.jetbrains.python.sdk.PythonSdkType
+
+private const val STUBS_SUFFIX = "-stubs"
+
+/**
+ * If [name] argument points to element in stub package,
+ * then [name] would be copied and `-stubs` suffix would be removed from the first component,
+ * otherwise [name] would be returned.
+ */
+fun convertStubToRuntimePackageName(name: QualifiedName): QualifiedName {
+  val top = name.firstComponent
+
+  if (top != null && top.endsWith(STUBS_SUFFIX)) {
+    return QualifiedName.fromComponents(name.components).apply { components[0] = components[0].dropLast(STUBS_SUFFIX.length) }
+  }
+
+  return name
+}
+
+/**
+ * Resolves [name] in corresponding stub package.
+ *
+ * Returns empty list if [context] disallow stubs,
+ * or language level is older than [LanguageLevel.PYTHON37],
+ * or [item] is not lib root.
+ */
+fun resolveModuleAtStubPackage(name: QualifiedName,
+                               item: PsiFileSystemItem,
+                               context: PyQualifiedNameResolveContext): List<PsiElement> {
+  if (!context.withoutStubs && name.componentCount > 0) {
+    val head = name.firstComponent!!
+
+    // prevent recursion and check that stub packages are allowed
+    if (!head.endsWith(STUBS_SUFFIX) && contextLanguageLevel(context).isAtLeast(LanguageLevel.PYTHON37)) {
+      val virtualFile = item.virtualFile
+
+      // 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 emptyList()
+}
+
+private fun contextLanguageLevel(context: PyQualifiedNameResolveContext): LanguageLevel {
+  context.foothold?.also { return LanguageLevel.forElement(it) }
+  context.footholdFile?.also { return LanguageLevel.forElement(it) }
+
+  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)
+}
+
+private fun pyi(element: PsiElement) = element is PyiFile || PyUtil.turnDirIntoInit(element) is PyiFile
@@ -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.resolveModuleAtStubPackage
 import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil
 import com.jetbrains.python.facet.PythonPathContributingFacet
 import com.jetbrains.python.psi.LanguageLevel
@@ -144,7 +145,7 @@ fun resolveModuleAt(name: QualifiedName, item: PsiFileSystemItem?, context: PyQu
                                                       !context.withPlainDirectories, context.withoutStubs, context.withoutForeign)
      PyUtil.filterTopPriorityResults(children.toTypedArray())
    }
-  }
+  } + resolveModuleAtStubPackage(name, item, context)
 }

 /**
@@ -24,6 +24,8 @@ import org.jetbrains.annotations.Nullable;

 import java.util.*;

+import static com.jetbrains.python.codeInsight.typing.PyPEP561.convertStubToRuntimePackageName;
+
 /**
 * @author yole
 */
@@ -207,21 +209,22 @@ public class QualifiedNameFinder {

    @Override
    public boolean visitRoot(@NotNull VirtualFile root, @Nullable Module module, @Nullable Sdk sdk, boolean isModuleSource) {
-      final List<String> result = pathToNameComponents(VfsUtilCore.getRelativePath(myVFile, root, '/'));
-      if (!result.isEmpty()) {
-        for (String component : result) {
+      final String relativePath = VfsUtilCore.getRelativePath(myVFile, root, '/');
+      final QualifiedName result = convertStubToRuntimePackageName(pathToQualifiedName(relativePath));
+      if (result.getComponentCount() != 0) {
+        for (String component : result.getComponents()) {
          if (!PyNames.isIdentifier(component)) {
            return true;
          }
        }
-        myResults.add(QualifiedName.fromComponents(result));
+        myResults.add(result);
      }
      return true;
    }

    @NotNull
-    private List<String> pathToNameComponents(@Nullable String relativePath) {
-      if (StringUtil.isEmpty(relativePath)) return Collections.emptyList();
+    private QualifiedName pathToQualifiedName(@Nullable String relativePath) {
+      if (StringUtil.isEmpty(relativePath)) return QualifiedName.fromComponents();

      final List<String> result = new ArrayList<>(StringUtil.split(relativePath, "/"));
      if (!result.isEmpty()) {
@@ -235,10 +238,16 @@ public class QualifiedNameFinder {
          result.remove(lastIndex);
        }

-        return result;
+        for (String component : result) {
+          if (component.contains(".")) {
+            return QualifiedName.fromComponents();
+          }
+        }
+
+        return QualifiedName.fromComponents(result);
      }

-      return Collections.emptyList();
+      return QualifiedName.fromComponents();
    }

    @NotNull
@@ -0,0 +1 @@
+def bar(a: str, b: str) -> str: ...
@@ -0,0 +1,2 @@
+def bar(a, b):
+    return a + b
@@ -0,0 +1,4 @@
+from pkg import foo
+
+foo.bar("a", "b")
+#    <ref>
@@ -0,0 +1 @@
+def bar(a: str, b: str) -> str: ...
@@ -0,0 +1,2 @@
+def bar(a, b):
+    return a + b
@@ -0,0 +1,4 @@
+import pkg.foo
+
+pkg.foo.bar("a", "b")
+#        <ref>
@@ -0,0 +1 @@
+def bar(a: str, b: str) -> str: ...
@@ -0,0 +1,2 @@
+def bar(a, b):
+    return a + b
@@ -0,0 +1,4 @@
+from pkg import foo
+
+foo.bar("a", "b")
+#    <ref>
@@ -15,6 +15,7 @@
 */
 package com.jetbrains.python;

+import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiReference;
 import com.intellij.psi.util.PsiTreeUtil;
@@ -357,7 +358,7 @@ public class Py3ResolveTest extends PyResolveTestCase {
  public void testDunderClassInDeclarationInsideFunction() {
    assertUnresolved();
  }
-  
+
  // PY-20864
  public void testTopLevelVariableAnnotationFromTyping() {
    runWithLanguageLevel(LanguageLevel.PYTHON36, () -> assertResolvesTo(PyElement.class, "List"));
@@ -654,4 +655,70 @@ public class Py3ResolveTest extends PyResolveTestCase {
  public void testRegexpAndFStringCombined() {
    runWithLanguageLevel(LanguageLevel.PYTHON36, () -> assertResolvesTo(PyTargetExpression.class, "foo"));
  }
+
+  // PY-30942
+  public void testStubPackage() {
+    myFixture.copyDirectoryToProject("resolve/" + getTestName(false), "");
+    final VirtualFile libDir = myFixture.findFileInTempDir("lib");
+    assertNotNull(libDir);
+
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON37,
+      () ->
+        runWithAdditionalClassEntryInSdkRoots(
+          libDir,
+          () -> {
+            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 testStubPackageFullyQName() {
+    myFixture.copyDirectoryToProject("resolve/" + getTestName(false), "");
+    final VirtualFile libDir = myFixture.findFileInTempDir("lib");
+    assertNotNull(libDir);
+
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON37,
+      () ->
+        runWithAdditionalClassEntryInSdkRoots(
+          libDir,
+          () -> {
+            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 testStubPackagePy36() {
+    myFixture.copyDirectoryToProject("resolve/StubPackage", "");
+    final VirtualFile libDir = myFixture.findFileInTempDir("lib");
+    assertNotNull(libDir);
+
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON36,
+      () ->
+        runWithAdditionalClassEntryInSdkRoots(
+          libDir,
+          () -> {
+            myFixture.configureByFile("main.py");
+
+            final PsiElement element = PyResolveTestCase.findReferenceByMarker(myFixture.getFile()).resolve();
+            assertInstanceOf(element, PyFunction.class);
+            assertEquals("foo.py", element.getContainingFile().getName());
+          }
+        )
+    );
+  }
 }