PY-20057 Support Type[Union[Class1, Class2]] annotation syntax

python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java

@@ -529,6 +529,12 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
           if (typeVar != null && !typeVar.isDefinition()) {
             return Ref.create(new PyGenericType(typeVar.getName(), typeVar.getBound(), true));
           }
+          // Represent Type[Union[str, int]] internally as Union[Type[str], Type[int]]
+          final PyUnionType unionType = as(type, PyUnionType.class);
+          if (unionType != null &&
+              unionType.getMembers().stream().allMatch(t -> t instanceof PyClassType && !((PyClassType)t).isDefinition())) {
+            return Ref.create(PyUnionType.union(ContainerUtil.map(unionType.getMembers(), t -> ((PyClassType)t).toClass())));
+          }
         }
         // Map Type[Something] with unsupported type parameter to Any, instead of generic type for the class "type"
         return Ref.create();

python/src/com/jetbrains/python/documentation/PyTypeModelBuilder.java

@@ -267,14 +267,19 @@ public class PyTypeModelBuilder {
     }
     else if (type instanceof PyUnionType && allowUnions) {
       final PyUnionType unionType = (PyUnionType)type;
+      final Collection<PyType> unionMembers = unionType.getMembers();
       if (type instanceof PyDynamicallyEvaluatedType || PyTypeChecker.isUnknown(type, false)) {
         result = new UnknownType(build(unionType.excludeNull(myContext), true));
       }
+      else if (unionMembers.stream().allMatch(t -> t instanceof PyClassType && ((PyClassType)t).isDefinition())) {
+        final List<TypeModel> instanceTypes = ContainerUtil.map(unionMembers, t -> build(((PyClassType)t).toInstance(), allowUnions));
+        result = new ClassObjectType(new OneOf(instanceTypes));
+      }
       else {
         result = Optional
           .ofNullable(getOptionalType(unionType))
-          .<PyTypeModelBuilder.TypeModel>map(optionalType -> new OptionalType(build(optionalType, true)))
-          .orElseGet(() -> new OneOf(Collections2.transform(unionType.getMembers(), t -> build(t, false))));
+          .<TypeModel>map(optionalType -> new OptionalType(build(optionalType, true)))
+          .orElseGet(() -> new OneOf(Collections2.transform(unionMembers, t -> build(t, false))));
       }
     }
     else if (type instanceof PyCallableType && !(type instanceof PyClassLikeType)) {

python/testData/inspections/PyTypeCheckerInspection/ClassObjectTypeWithUnion.py

@@ -0,0 +1,21 @@
+from typing import Type, Union
+
+class MyClass:
+    pass
+
+def expects_myclass_or_str1(x: Type[Union[MyClass, str]]):
+    pass
+
+expects_myclass_or_str1(MyClass)
+expects_myclass_or_str1(str)
+expects_myclass_or_str1(<warning descr="Expected type 'Type[Union[MyClass, str]]', got 'Type[int]' instead">int</warning>)
+expects_myclass_or_str1(<warning descr="Expected type 'Type[Union[MyClass, str]]', got 'int' instead">42</warning>)
+
+
+def expects_myclass_or_str2(x: Union[Type[MyClass], Type[str]]):
+    pass
+
+expects_myclass_or_str2(MyClass)
+expects_myclass_or_str2(str)
+expects_myclass_or_str2(<warning descr="Expected type 'Type[Union[MyClass, str]]', got 'Type[int]' instead">int</warning>)
+expects_myclass_or_str2(<warning descr="Expected type 'Type[Union[MyClass, str]]', got 'int' instead">42</warning>)

python/testSrc/com/jetbrains/python/PyTypingTest.java

@@ -858,6 +858,15 @@ public class PyTypingTest extends PyTestCase {
            "def f(x: Tuple[Type[42], Type[], Type[unresolved]]):\n" +
            "    expr = x");
   }
+  
+  // PY-20057
+  public void testUnionOfClassObjectTypes() {
+    doTest("Type[Union[int, str]]",
+           "from typing import Type, Union\n" +
+           "\n" +
+           "def f(x: Type[Union[int, str]]):\n" +
+           "    expr = x");
+  }
 
   private void doTestNoInjectedText(@NotNull String text) {
     myFixture.configureByText(PythonFileType.INSTANCE, text);

python/testSrc/com/jetbrains/python/inspections/Py3TypeCheckerInspectionTest.java

@@ -208,4 +208,9 @@ public class Py3TypeCheckerInspectionTest extends PyTestCase {
   public void testTypeAndClassObjectTypesCompatibility() {
     doTest();
   }
+
+  // PY-20057
+  public void testClassObjectTypeWithUnion() {
+    doTest();
+  }
 }