PY-75291: Implemented auto-import for PEP 695 type aliases

GitOrigin-RevId: c15a6387785813688358e9695186a72dbdf35744
2026-09-27 10:03:11 +07:00 · 2025-01-23 15:46:39 +00:00
parent 3076cfe8e2
commit 02f5e3f522
11 changed files with 67 additions and 12 deletions
@@ -38,6 +38,7 @@
    <stubIndex implementation="com.jetbrains.python.psi.stubs.PyDecoratorStubIndex"/>
    <stubIndex implementation="com.jetbrains.python.psi.stubs.PyClassAttributesIndex"/>
    <stubIndex implementation="com.jetbrains.python.psi.stubs.PyExportedModuleAttributeIndex"/>
+    <stubIndex implementation="com.jetbrains.python.psi.stubs.PyTypeAliasNameIndex"/>
    <fileBasedIndex implementation="com.jetbrains.python.psi.stubs.PyModuleNameIndex"/>
    <fileBasedIndex implementation="com.jetbrains.python.psi.stubs.PySetuptoolsNamespaceIndex"/>

@@ -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<PsiNamedElement> 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;
  }
 }
@@ -60,7 +60,7 @@ public class PyFileElementType extends IStubFileElementType<PyFileStub> {
  @Override
  public int getStubVersion() {
    // Don't forget to update versions of indexes that use the updated stub-based elements
-    return 95;
+    return 96;
  }

  @Override
@@ -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<PyTypeAli
                                            parentStub, getStubElementType(), PyVersionSpecificStubBaseKt.evaluateVersionsForElement(psi));
  }

+  @Override
+  public void indexStub(@NotNull PyTypeAliasStatementStub stub, @NotNull IndexSink sink) {
+    String name = stub.getName();
+    if (name != null && PyUtil.getInitialUnderscores(name) == 0 && stub.getParentStub() instanceof PyFileStub) {
+      sink.occurrence(PyTypeAliasNameIndex.KEY, name);
+    }
+  }
+
  @Override
  public void serialize(@NotNull PyTypeAliasStatementStub stub, @NotNull StubOutputStream dataStream) throws IOException {
    dataStream.writeName(stub.getName());
@@ -0,0 +1,26 @@
+// Copyright 2000-2025 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.
+package com.jetbrains.python.psi.stubs;
+
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.search.GlobalSearchScope;
+import com.intellij.psi.stubs.StringStubIndexExtension;
+import com.intellij.psi.stubs.StubIndex;
+import com.intellij.psi.stubs.StubIndexKey;
+import com.jetbrains.python.psi.PyTypeAliasStatement;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Collection;
+
+
+public final class PyTypeAliasNameIndex extends StringStubIndexExtension<PyTypeAliasStatement> {
+  public static final StubIndexKey<String, PyTypeAliasStatement> KEY = StubIndexKey.createIndexKey("Py.TypeAliasName");
+
+  @Override
+  public @NotNull StubIndexKey<String, PyTypeAliasStatement> getKey() {
+    return KEY;
+  }
+
+  public static @NotNull Collection<PyTypeAliasStatement> find(String name, Project project, GlobalSearchScope scope) {
+    return StubIndex.getElements(KEY, name, project, scope, PyTypeAliasStatement.class);
+  }
+}
@@ -0,0 +1 @@
+type ExampleType = int
@@ -0,0 +1 @@
+x: <error descr="Unresolved reference 'ExampleType'">ExampleTy<caret>pe</error>
@@ -0,0 +1,3 @@
+from lib import ExampleType
+
+x: ExampleType
@@ -0,0 +1,2 @@
+type PublicType = int
+type _PrivateType = int
@@ -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();
@@ -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<String> candidates = ContainerUtil.map(fix.getCandidates(), c -> c.getPresentableText());
				`@@ -0,0 +1 @@`
				`x: <error descr="Unresolved reference 'ExampleType'">ExampleTy<caret>pe</error>`