From 02f5e3f522efae780d33c8b494b91650ccdb038d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alex=20Gr=C3=B6nholm?= Date: Mon, 20 Jan 2025 16:21:37 +0200 Subject: [PATCH] PY-75291: Implemented auto-import for PEP 695 type aliases GitOrigin-RevId: c15a6387785813688358e9695186a72dbdf35744 --- .../resources/META-INF/PythonPsiImpl.xml | 1 + .../imports/PyImportCollector.java | 15 ++++++----- .../python/psi/PyFileElementType.java | 2 +- .../PyTypeAliasStatementElementType.java | 16 ++++++++---- .../psi/stubs/PyTypeAliasNameIndex.java | 26 +++++++++++++++++++ .../typeAliasStatement/lib.py | 1 + .../typeAliasStatement/main.py | 1 + .../typeAliasStatement/main_after.py | 3 +++ python/testData/stubs/TypeAliasNameIndex.py | 2 ++ .../com/jetbrains/python/PyStubsTest.java | 8 ++++++ .../quickFixes/PyAddImportQuickFixTest.java | 4 +++ 11 files changed, 67 insertions(+), 12 deletions(-) create mode 100644 python/python-psi-impl/src/com/jetbrains/python/psi/stubs/PyTypeAliasNameIndex.java create mode 100644 python/testData/quickFixes/PyAddImportQuickFixTest/typeAliasStatement/lib.py create mode 100644 python/testData/quickFixes/PyAddImportQuickFixTest/typeAliasStatement/main.py create mode 100644 python/testData/quickFixes/PyAddImportQuickFixTest/typeAliasStatement/main_after.py create mode 100644 python/testData/stubs/TypeAliasNameIndex.py diff --git a/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml b/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml index 53863cd5a328..d835c2c9c755 100644 --- a/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml +++ b/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml @@ -38,6 +38,7 @@ + diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/PyImportCollector.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/PyImportCollector.java index 9482c261f0be..dfaf52747e00 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/PyImportCollector.java +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/PyImportCollector.java @@ -10,15 +10,14 @@ import com.intellij.psi.util.PsiTreeUtil; import com.intellij.psi.util.QualifiedName; import com.jetbrains.python.PyNames; import com.jetbrains.python.codeInsight.PyCodeInsightSettings; +import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider; import com.jetbrains.python.inspections.unresolvedReference.PyCommonImportAliasesKt; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyFileImpl; import com.jetbrains.python.psi.resolve.QualifiedNameFinder; import com.jetbrains.python.psi.search.PySearchUtilBase; -import com.jetbrains.python.psi.stubs.PyClassNameIndex; -import com.jetbrains.python.psi.stubs.PyFunctionNameIndex; -import com.jetbrains.python.psi.stubs.PyModuleNameIndex; -import com.jetbrains.python.psi.stubs.PyVariableNameIndex; +import com.jetbrains.python.psi.stubs.*; +import com.jetbrains.python.psi.types.TypeEvalContext; import org.jetbrains.annotations.NotNull; import java.util.*; @@ -106,12 +105,16 @@ public class PyImportCollector { private void addSymbolImportCandidates(PsiFile existingImportFile) { Project project = myNode.getProject(); GlobalSearchScope scope = PySearchUtilBase.defaultSuggestionScope(myNode); + TypeEvalContext context = TypeEvalContext.codeAnalysis(project, myNode.getContainingFile()); List symbols = new ArrayList<>(PyClassNameIndex.find(myRefText, project, scope)); if (!isQualifier()) { symbols.addAll(PyFunctionNameIndex.find(myRefText, project, scope)); } symbols.addAll(PyVariableNameIndex.find(myRefText, project, scope)); + if (PyTypingTypeProvider.isInsideTypeHint(myNode, context)) { + symbols.addAll(PyTypeAliasNameIndex.find(myRefText, project, scope)); + } if (isPossibleModuleReference()) { symbols.addAll(findImportableModules(myRefText, false, scope)); String packageQName = PyCommonImportAliasesKt.PY_COMMON_IMPORT_ALIASES.get(myRefText); @@ -209,7 +212,7 @@ public class PyImportCollector { if (symbol instanceof PyClass || symbol instanceof PyFunction) { return PyUtil.isTopLevel(symbol); } - // only top-level target expressions are included in VariableNameIndex - return symbol instanceof PyTargetExpression; + // only top-level target expressions and type aliases are included in VariableNameIndex and TypeAliasNameIndex, respectively + return symbol instanceof PyTargetExpression || symbol instanceof PyTypeAliasStatement; } } diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/PyFileElementType.java b/python/python-psi-impl/src/com/jetbrains/python/psi/PyFileElementType.java index 7d6359196385..0859cea62523 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/PyFileElementType.java +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/PyFileElementType.java @@ -60,7 +60,7 @@ public class PyFileElementType extends IStubFileElementType { @Override public int getStubVersion() { // Don't forget to update versions of indexes that use the updated stub-based elements - return 95; + return 96; } @Override diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyTypeAliasStatementElementType.java b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyTypeAliasStatementElementType.java index af6e980931c8..61d55e0586fe 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyTypeAliasStatementElementType.java +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyTypeAliasStatementElementType.java @@ -4,15 +4,13 @@ import com.google.common.collect.RangeSet; import com.intellij.lang.ASTNode; import com.intellij.openapi.util.Version; import com.intellij.psi.PsiElement; -import com.intellij.psi.stubs.IStubElementType; -import com.intellij.psi.stubs.StubElement; -import com.intellij.psi.stubs.StubInputStream; -import com.intellij.psi.stubs.StubOutputStream; +import com.intellij.psi.stubs.*; import com.jetbrains.python.PyStubElementTypes; import com.jetbrains.python.psi.PyStubElementType; import com.jetbrains.python.psi.PyTypeAliasStatement; +import com.jetbrains.python.psi.PyUtil; import com.jetbrains.python.psi.impl.PyTypeAliasStatementImpl; -import com.jetbrains.python.psi.stubs.PyTypeAliasStatementStub; +import com.jetbrains.python.psi.stubs.*; import org.jetbrains.annotations.NotNull; import java.io.IOException; @@ -39,6 +37,14 @@ public class PyTypeAliasStatementElementType extends PyStubElementType { + public static final StubIndexKey KEY = StubIndexKey.createIndexKey("Py.TypeAliasName"); + + @Override + public @NotNull StubIndexKey getKey() { + return KEY; + } + + public static @NotNull Collection find(String name, Project project, GlobalSearchScope scope) { + return StubIndex.getElements(KEY, name, project, scope, PyTypeAliasStatement.class); + } +} diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/typeAliasStatement/lib.py b/python/testData/quickFixes/PyAddImportQuickFixTest/typeAliasStatement/lib.py new file mode 100644 index 000000000000..d32274132b0e --- /dev/null +++ b/python/testData/quickFixes/PyAddImportQuickFixTest/typeAliasStatement/lib.py @@ -0,0 +1 @@ +type ExampleType = int \ No newline at end of file diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/typeAliasStatement/main.py b/python/testData/quickFixes/PyAddImportQuickFixTest/typeAliasStatement/main.py new file mode 100644 index 000000000000..1472ff54e7ab --- /dev/null +++ b/python/testData/quickFixes/PyAddImportQuickFixTest/typeAliasStatement/main.py @@ -0,0 +1 @@ +x: ExampleType diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/typeAliasStatement/main_after.py b/python/testData/quickFixes/PyAddImportQuickFixTest/typeAliasStatement/main_after.py new file mode 100644 index 000000000000..069ff3eb599b --- /dev/null +++ b/python/testData/quickFixes/PyAddImportQuickFixTest/typeAliasStatement/main_after.py @@ -0,0 +1,3 @@ +from lib import ExampleType + +x: ExampleType diff --git a/python/testData/stubs/TypeAliasNameIndex.py b/python/testData/stubs/TypeAliasNameIndex.py new file mode 100644 index 000000000000..1f03bac1cfec --- /dev/null +++ b/python/testData/stubs/TypeAliasNameIndex.py @@ -0,0 +1,2 @@ +type PublicType = int +type _PrivateType = int \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyStubsTest.java b/python/testSrc/com/jetbrains/python/PyStubsTest.java index c1be306a2bff..0d5cc2cd538f 100644 --- a/python/testSrc/com/jetbrains/python/PyStubsTest.java +++ b/python/testSrc/com/jetbrains/python/PyStubsTest.java @@ -859,6 +859,14 @@ public class PyStubsTest extends PyTestCase { assertNotParsed(file); } + // PY-75291 + public void testTypeAliasNameIndex() { + getTestFile(); + GlobalSearchScope scope = GlobalSearchScope.allScope(myFixture.getProject()); + assertEquals(1, PyTypeAliasNameIndex.find("PublicType", myFixture.getProject(), scope).size()); + assertEquals(0, PyTypeAliasNameIndex.find("_PrivateType", myFixture.getProject(), scope).size()); + } + // PY-18866 public void testUnresolvedTypingSymbol() { final PyFile file = getTestFile(); diff --git a/python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java b/python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java index ec7e7f9276ca..368d15744b3f 100644 --- a/python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java +++ b/python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java @@ -425,6 +425,10 @@ public class PyAddImportQuickFixTest extends PyQuickFixTestCase { doMultiFileAutoImportTest("Import 'decimal.Decimal'"); } + public void testTypeAliasStatement() { + doMultiFileAutoImportTest("Import 'lib.ExampleType'"); + } + private void doTestProposedImportsOrdering(String @NotNull ... expected) { doMultiFileAutoImportTest("Import", fix -> { final List candidates = ContainerUtil.map(fix.getCandidates(), c -> c.getPresentableText());