[python] PY-80238 add completion for TypeAlias in type annotations

GitOrigin-RevId: e5f66e1fcf75b3eb42d06c5700d6b3b9de325a03

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;
   }

python/testData/completion/className/newTypeAliasCompletion/a.py

@@ -0,0 +1 @@
+x: NewStyle<caret>

python/testData/completion/className/newTypeAliasCompletion/type_aliases.py

@@ -0,0 +1,3 @@
+type NewStyleString = str
+type NewStyleNumber = int | float
+type NewStyleUnion = str | int | None

python/testData/completion/className/oldStyleTypeAliasCompletion/a.py

@@ -0,0 +1 @@
+x : OldStyle<caret>

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

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();