From 7863e90505f391f628b7d0bdd2bc854f0ec57e89 Mon Sep 17 00:00:00 2001 From: Semyon Proshev Date: Mon, 8 Oct 2018 15:53:30 +0300 Subject: [PATCH] Add resolve to stub packages (PY-30942) --- .../python/codeInsight/typing/PyPEP561.kt | 73 +++++++++++++++++++ .../python/psi/resolve/PyResolveImportUtil.kt | 3 +- .../psi/resolve/QualifiedNameFinder.java | 25 +++++-- .../StubPackage/lib/pkg-stubs/__init__.pyi | 0 .../resolve/StubPackage/lib/pkg-stubs/foo.pyi | 1 + .../resolve/StubPackage/lib/pkg/__init__.py | 0 .../resolve/StubPackage/lib/pkg/foo.py | 2 + python/testData/resolve/StubPackage/main.py | 4 + .../lib/pkg-stubs/__init__.pyi | 0 .../lib/pkg-stubs/foo.pyi | 1 + .../StubPackageFullyQName/lib/pkg/__init__.py | 0 .../StubPackageFullyQName/lib/pkg/foo.py | 2 + .../resolve/StubPackageFullyQName/main.py | 4 + .../lib/pkg-stubs/__init__.pyi | 0 .../StubPackagePy36/lib/pkg-stubs/foo.pyi | 1 + .../StubPackagePy36/lib/pkg/__init__.py | 0 .../resolve/StubPackagePy36/lib/pkg/foo.py | 2 + .../testData/resolve/StubPackagePy36/main.py | 4 + .../com/jetbrains/python/Py3ResolveTest.java | 69 +++++++++++++++++- 19 files changed, 181 insertions(+), 10 deletions(-) create mode 100644 python/src/com/jetbrains/python/codeInsight/typing/PyPEP561.kt create mode 100644 python/testData/resolve/StubPackage/lib/pkg-stubs/__init__.pyi create mode 100644 python/testData/resolve/StubPackage/lib/pkg-stubs/foo.pyi create mode 100644 python/testData/resolve/StubPackage/lib/pkg/__init__.py create mode 100644 python/testData/resolve/StubPackage/lib/pkg/foo.py create mode 100644 python/testData/resolve/StubPackage/main.py create mode 100644 python/testData/resolve/StubPackageFullyQName/lib/pkg-stubs/__init__.pyi create mode 100644 python/testData/resolve/StubPackageFullyQName/lib/pkg-stubs/foo.pyi create mode 100644 python/testData/resolve/StubPackageFullyQName/lib/pkg/__init__.py create mode 100644 python/testData/resolve/StubPackageFullyQName/lib/pkg/foo.py create mode 100644 python/testData/resolve/StubPackageFullyQName/main.py create mode 100644 python/testData/resolve/StubPackagePy36/lib/pkg-stubs/__init__.pyi create mode 100644 python/testData/resolve/StubPackagePy36/lib/pkg-stubs/foo.pyi create mode 100644 python/testData/resolve/StubPackagePy36/lib/pkg/__init__.py create mode 100644 python/testData/resolve/StubPackagePy36/lib/pkg/foo.py create mode 100644 python/testData/resolve/StubPackagePy36/main.py diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyPEP561.kt b/python/src/com/jetbrains/python/codeInsight/typing/PyPEP561.kt new file mode 100644 index 000000000000..c83e3b318742 --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/typing/PyPEP561.kt @@ -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 { + 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 \ No newline at end of file diff --git a/python/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt b/python/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt index ec52ae5a8522..25faf49536a5 100644 --- a/python/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt +++ b/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.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) } /** diff --git a/python/src/com/jetbrains/python/psi/resolve/QualifiedNameFinder.java b/python/src/com/jetbrains/python/psi/resolve/QualifiedNameFinder.java index 892d117eeb20..ea0cecd76fb5 100644 --- a/python/src/com/jetbrains/python/psi/resolve/QualifiedNameFinder.java +++ b/python/src/com/jetbrains/python/psi/resolve/QualifiedNameFinder.java @@ -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 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 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 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 diff --git a/python/testData/resolve/StubPackage/lib/pkg-stubs/__init__.pyi b/python/testData/resolve/StubPackage/lib/pkg-stubs/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/StubPackage/lib/pkg-stubs/foo.pyi b/python/testData/resolve/StubPackage/lib/pkg-stubs/foo.pyi new file mode 100644 index 000000000000..36a173b06914 --- /dev/null +++ b/python/testData/resolve/StubPackage/lib/pkg-stubs/foo.pyi @@ -0,0 +1 @@ +def bar(a: str, b: str) -> str: ... \ No newline at end of file diff --git a/python/testData/resolve/StubPackage/lib/pkg/__init__.py b/python/testData/resolve/StubPackage/lib/pkg/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/StubPackage/lib/pkg/foo.py b/python/testData/resolve/StubPackage/lib/pkg/foo.py new file mode 100644 index 000000000000..85e2338108b8 --- /dev/null +++ b/python/testData/resolve/StubPackage/lib/pkg/foo.py @@ -0,0 +1,2 @@ +def bar(a, b): + return a + b \ No newline at end of file diff --git a/python/testData/resolve/StubPackage/main.py b/python/testData/resolve/StubPackage/main.py new file mode 100644 index 000000000000..60e25e4ed882 --- /dev/null +++ b/python/testData/resolve/StubPackage/main.py @@ -0,0 +1,4 @@ +from pkg import foo + +foo.bar("a", "b") +# \ No newline at end of file diff --git a/python/testData/resolve/StubPackageFullyQName/lib/pkg-stubs/__init__.pyi b/python/testData/resolve/StubPackageFullyQName/lib/pkg-stubs/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/StubPackageFullyQName/lib/pkg-stubs/foo.pyi b/python/testData/resolve/StubPackageFullyQName/lib/pkg-stubs/foo.pyi new file mode 100644 index 000000000000..36a173b06914 --- /dev/null +++ b/python/testData/resolve/StubPackageFullyQName/lib/pkg-stubs/foo.pyi @@ -0,0 +1 @@ +def bar(a: str, b: str) -> str: ... \ No newline at end of file diff --git a/python/testData/resolve/StubPackageFullyQName/lib/pkg/__init__.py b/python/testData/resolve/StubPackageFullyQName/lib/pkg/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/StubPackageFullyQName/lib/pkg/foo.py b/python/testData/resolve/StubPackageFullyQName/lib/pkg/foo.py new file mode 100644 index 000000000000..85e2338108b8 --- /dev/null +++ b/python/testData/resolve/StubPackageFullyQName/lib/pkg/foo.py @@ -0,0 +1,2 @@ +def bar(a, b): + return a + b \ No newline at end of file diff --git a/python/testData/resolve/StubPackageFullyQName/main.py b/python/testData/resolve/StubPackageFullyQName/main.py new file mode 100644 index 000000000000..a54bb87a6e7e --- /dev/null +++ b/python/testData/resolve/StubPackageFullyQName/main.py @@ -0,0 +1,4 @@ +import pkg.foo + +pkg.foo.bar("a", "b") +# \ No newline at end of file diff --git a/python/testData/resolve/StubPackagePy36/lib/pkg-stubs/__init__.pyi b/python/testData/resolve/StubPackagePy36/lib/pkg-stubs/__init__.pyi new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/StubPackagePy36/lib/pkg-stubs/foo.pyi b/python/testData/resolve/StubPackagePy36/lib/pkg-stubs/foo.pyi new file mode 100644 index 000000000000..36a173b06914 --- /dev/null +++ b/python/testData/resolve/StubPackagePy36/lib/pkg-stubs/foo.pyi @@ -0,0 +1 @@ +def bar(a: str, b: str) -> str: ... \ No newline at end of file diff --git a/python/testData/resolve/StubPackagePy36/lib/pkg/__init__.py b/python/testData/resolve/StubPackagePy36/lib/pkg/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/StubPackagePy36/lib/pkg/foo.py b/python/testData/resolve/StubPackagePy36/lib/pkg/foo.py new file mode 100644 index 000000000000..85e2338108b8 --- /dev/null +++ b/python/testData/resolve/StubPackagePy36/lib/pkg/foo.py @@ -0,0 +1,2 @@ +def bar(a, b): + return a + b \ No newline at end of file diff --git a/python/testData/resolve/StubPackagePy36/main.py b/python/testData/resolve/StubPackagePy36/main.py new file mode 100644 index 000000000000..60e25e4ed882 --- /dev/null +++ b/python/testData/resolve/StubPackagePy36/main.py @@ -0,0 +1,4 @@ +from pkg import foo + +foo.bar("a", "b") +# \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/Py3ResolveTest.java b/python/testSrc/com/jetbrains/python/Py3ResolveTest.java index 8a94671a3f40..a3c0691b6eca 100644 --- a/python/testSrc/com/jetbrains/python/Py3ResolveTest.java +++ b/python/testSrc/com/jetbrains/python/Py3ResolveTest.java @@ -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()); + } + ) + ); + } }