diff --git a/crates/red_knot_python_semantic/resources/mdtest/type_of/basic.md b/crates/red_knot_python_semantic/resources/mdtest/type_of/basic.md index c1b9a28e3ec49..15326a8f65a45 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/type_of/basic.md +++ b/crates/red_knot_python_semantic/resources/mdtest/type_of/basic.md @@ -87,3 +87,31 @@ reveal_type(f()) # revealed: @Todo(unsupported type[X] special form) ```py path=a/b.py class C: ... ``` + +## Union of classes + +```py +class BasicUser: ... +class ProUser: ... + +class A: + class B: + class C: ... + +def get_user() -> type[BasicUser | ProUser | A.B.C]: + return BasicUser + +# revealed: type[BasicUser] | type[ProUser] | type[C] +reveal_type(get_user()) +``` + +## Illegal parameters + +```py +class A: ... +class B: ... + +# error: [invalid-type-form] +def get_user() -> type[A, B]: + return A +``` diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index e716864f3a029..0e2fe388f6ae6 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -4663,7 +4663,28 @@ impl<'db> TypeInferenceBuilder<'db> { todo_type!("unsupported type[X] special form") } } - // TODO: unions, subscripts, etc. + ast::Expr::BinOp(binary) if binary.op == ast::Operator::BitOr => { + let union_ty = UnionType::from_elements( + self.db, + [ + self.infer_subclass_of_type_expression(&binary.left), + self.infer_subclass_of_type_expression(&binary.right), + ], + ); + self.store_expression_type(slice, union_ty); + + union_ty + } + ast::Expr::Tuple(_) => { + self.infer_type_expression(slice); + self.diagnostics.add( + slice.into(), + "invalid-type-form", + format_args!("type[...] must have exactly one type argument"), + ); + Type::Unknown + } + // TODO: subscripts, etc. _ => { self.infer_type_expression(slice); todo_type!("unsupported type[X] special form")