From 57904e8c2d7697f9fc25f15e02bdb4634405b0df Mon Sep 17 00:00:00 2001
From: Azim Akhmadjonov <azimjon.akhmadjonov@jetbrains.com>
Date: Mon, 3 Nov 2025 12:47:39 +0100
Subject: [PATCH] [python] PY-80238 add completion for TypeAlias in type
 annotations

GitOrigin-RevId: e5f66e1fcf75b3eb42d06c5700d6b3b9de325a03
---
 .../PyClassNameCompletionContributor.java      |  9 ++++++++-
 .../className/newTypeAliasCompletion/a.py      |  1 +
 .../newTypeAliasCompletion/type_aliases.py     |  3 +++
 .../className/oldStyleTypeAliasCompletion/a.py |  1 +
 .../type_aliases.py                            |  5 +++++
 .../python/PyClassNameCompletionTest.java      | 18 ++++++++++++++++++
 6 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 python/testData/completion/className/newTypeAliasCompletion/a.py
 create mode 100644 python/testData/completion/className/newTypeAliasCompletion/type_aliases.py
 create mode 100644 python/testData/completion/className/oldStyleTypeAliasCompletion/a.py
 create mode 100644 python/testData/completion/className/oldStyleTypeAliasCompletion/type_aliases.py

diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyClassNameCompletionContributor.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyClassNameCompletionContributor.java
index 6f877e6d3993..29dba731793d 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyClassNameCompletionContributor.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyClassNameCompletionContributor.java
@@ -206,7 +206,7 @@ public final class PyClassNameCompletionContributor extends PyImportableNameComp
       // Not all names from typing.py are defined as classes
       return (
         definition instanceof PyClass ||
-        definition instanceof PyTypeAliasStatement ||
+        isSuitableTypeAlias(definition, context) ||
         ArrayUtil.contains(fqn.getFirstComponent(), "typing", "typing_extensions")
       );
     }
@@ -216,6 +216,13 @@ public final class PyClassNameCompletionContributor extends PyImportableNameComp
     return true;
   }
 
+  private static boolean isSuitableTypeAlias(@NotNull PyElement element, @NotNull TypeEvalContext context) {
+    if (element instanceof PyTargetExpression && element.getParent() instanceof PyAssignmentStatement assignmentStatement) {
+      if (PyTypingTypeProvider.isExplicitTypeAlias(assignmentStatement, context)) return true;
+    }
+    return element instanceof PyTypeAliasStatement;
+  }
+
   private static boolean isInsideErrorElement(@NotNull PyReferenceExpression referenceExpression) {
     return PsiTreeUtil.getParentOfType(referenceExpression, PsiErrorElement.class) != null;
   }
diff --git a/python/testData/completion/className/newTypeAliasCompletion/a.py b/python/testData/completion/className/newTypeAliasCompletion/a.py
new file mode 100644
index 000000000000..96336a839de1
--- /dev/null
+++ b/python/testData/completion/className/newTypeAliasCompletion/a.py
@@ -0,0 +1 @@
+x: NewStyle<caret>
\ No newline at end of file
diff --git a/python/testData/completion/className/newTypeAliasCompletion/type_aliases.py b/python/testData/completion/className/newTypeAliasCompletion/type_aliases.py
new file mode 100644
index 000000000000..4f62d97f0894
--- /dev/null
+++ b/python/testData/completion/className/newTypeAliasCompletion/type_aliases.py
@@ -0,0 +1,3 @@
+type NewStyleString = str
+type NewStyleNumber = int | float
+type NewStyleUnion = str | int | None
\ No newline at end of file
diff --git a/python/testData/completion/className/oldStyleTypeAliasCompletion/a.py b/python/testData/completion/className/oldStyleTypeAliasCompletion/a.py
new file mode 100644
index 000000000000..4693bba08e60
--- /dev/null
+++ b/python/testData/completion/className/oldStyleTypeAliasCompletion/a.py
@@ -0,0 +1 @@
+x : OldStyle<caret>
\ No newline at end of file
diff --git a/python/testData/completion/className/oldStyleTypeAliasCompletion/type_aliases.py b/python/testData/completion/className/oldStyleTypeAliasCompletion/type_aliases.py
new file mode 100644
index 000000000000..39665ad8cedd
--- /dev/null
+++ b/python/testData/completion/className/oldStyleTypeAliasCompletion/type_aliases.py
@@ -0,0 +1,5 @@
+from typing import TypeAlias
+
+OldStyleString: TypeAlias = str
+OldStyleNumber: TypeAlias = int | float
+OldStyleUnion: TypeAlias = str | int | None
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PyClassNameCompletionTest.java b/python/testSrc/com/jetbrains/python/PyClassNameCompletionTest.java
index a2db21694ae5..85448c5d2bbe 100644
--- a/python/testSrc/com/jetbrains/python/PyClassNameCompletionTest.java
+++ b/python/testSrc/com/jetbrains/python/PyClassNameCompletionTest.java
@@ -199,6 +199,24 @@ public class PyClassNameCompletionTest extends PyTestCase {
     });
   }
 
+  // PY-80238
+  public void testOldStyleTypeAliasCompletion() {
+    myFixture.copyDirectoryToProject(getTestName(true), "");
+    myFixture.configureByFile("a.py");
+    myFixture.complete(CompletionType.BASIC, 1);
+    List<String> variants = myFixture.getLookupElementStrings();
+    assertContainsElements(variants, "OldStyleString", "OldStyleNumber", "OldStyleUnion");
+  }
+
+  // PY-80238
+  public void testNewTypeAliasCompletion() {
+    myFixture.copyDirectoryToProject(getTestName(true), "");
+    myFixture.configureByFile("a.py");
+    myFixture.complete(CompletionType.BASIC, 1);
+    List<String> variants = myFixture.getLookupElementStrings();
+    assertContainsElements(variants, "NewStyleString", "NewStyleNumber", "NewStyleUnion");
+  }
+
   @Nullable
   private static String getElementQualifiedName(@NotNull LookupElement le) {
     PsiElement element = le.getPsiElement();