diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java index 9ec2fe8ac819..6ab2f86239f7 100644 --- a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java +++ b/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(); diff --git a/python/src/com/jetbrains/python/documentation/PyTypeModelBuilder.java b/python/src/com/jetbrains/python/documentation/PyTypeModelBuilder.java index 017b812fc5c3..80a36ad769e4 100644 --- a/python/src/com/jetbrains/python/documentation/PyTypeModelBuilder.java +++ b/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 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 instanceTypes = ContainerUtil.map(unionMembers, t -> build(((PyClassType)t).toInstance(), allowUnions)); + result = new ClassObjectType(new OneOf(instanceTypes)); + } else { result = Optional .ofNullable(getOptionalType(unionType)) - .map(optionalType -> new OptionalType(build(optionalType, true))) - .orElseGet(() -> new OneOf(Collections2.transform(unionType.getMembers(), t -> build(t, false)))); + .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)) { diff --git a/python/testData/inspections/PyTypeCheckerInspection/ClassObjectTypeWithUnion.py b/python/testData/inspections/PyTypeCheckerInspection/ClassObjectTypeWithUnion.py new file mode 100644 index 000000000000..4e22f3045773 --- /dev/null +++ b/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(int) +expects_myclass_or_str1(42) + + +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(int) +expects_myclass_or_str2(42) diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java index 84ac1a202199..9342dd00adfe 100644 --- a/python/testSrc/com/jetbrains/python/PyTypingTest.java +++ b/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); diff --git a/python/testSrc/com/jetbrains/python/inspections/Py3TypeCheckerInspectionTest.java b/python/testSrc/com/jetbrains/python/inspections/Py3TypeCheckerInspectionTest.java index 33698bd3374e..fd89036b5c53 100644 --- a/python/testSrc/com/jetbrains/python/inspections/Py3TypeCheckerInspectionTest.java +++ b/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(); + } }