[red-knot] Understanding type[Union[A, B]] (#14858)

Co-authored-by: Alex Waygood <[email protected]>

crates/red_knot_python_semantic/resources/mdtest/type_of/basic.md

@@ -88,7 +88,7 @@ reveal_type(f())  # revealed: @Todo(unsupported type[X] special form)
 class C: ...
 ```
 
-## Union of classes
+## New-style union of classes
 
 ```py
 class BasicUser: ...
@@ -105,6 +105,41 @@ def get_user() -> type[BasicUser | ProUser | A.B.C]:
 reveal_type(get_user())
 ```
 
+## Old-style union of classes
+
+```py
+from typing import Union
+
+class BasicUser: ...
+class ProUser: ...
+
+class A:
+    class B:
+        class C: ...
+
+def f(a: type[Union[BasicUser, ProUser, A.B.C]], b: type[Union[str]], c: type[Union[BasicUser, Union[ProUser, A.B.C]]]):
+    reveal_type(a)  # revealed: type[BasicUser] | type[ProUser] | type[C]
+    reveal_type(b)  # revealed: type[str]
+    reveal_type(c)  # revealed: type[BasicUser] | type[ProUser] | type[C]
+```
+
+## New-style and old-style unions in combination
+
+```py
+from typing import Union
+
+class BasicUser: ...
+class ProUser: ...
+
+class A:
+    class B:
+        class C: ...
+
+def f(a: type[BasicUser | Union[ProUser, A.B.C]], b: type[Union[BasicUser | Union[ProUser, A.B.C | str]]]):
+    reveal_type(a)  # revealed: type[BasicUser] | type[ProUser] | type[C]
+    reveal_type(b)  # revealed: type[BasicUser] | type[ProUser] | type[C] | type[str]
+```
+
 ## Illegal parameters
 
 ```py

crates/red_knot_python_semantic/src/types/infer.rs

@@ -4684,6 +4684,33 @@ impl<'db> TypeInferenceBuilder<'db> {
                 );
                 Type::Unknown
             }
+            ast::Expr::Subscript(ast::ExprSubscript {
+                value,
+                slice: parameters,
+                ..
+            }) => {
+                let parameters_ty = match self.infer_expression(value) {
+                    Type::KnownInstance(KnownInstanceType::Union) => match &**parameters {
+                        ast::Expr::Tuple(tuple) => {
+                            let ty = UnionType::from_elements(
+                                self.db,
+                                tuple
+                                    .iter()
+                                    .map(|element| self.infer_subclass_of_type_expression(element)),
+                            );
+                            self.store_expression_type(parameters, ty);
+                            ty
+                        }
+                        _ => self.infer_subclass_of_type_expression(parameters),
+                    },
+                    _ => {
+                        self.infer_type_expression(parameters);
+                        todo_type!("unsupported nested subscript in type[X]")
+                    }
+                };
+                self.store_expression_type(slice, parameters_ty);
+                parameters_ty
+            }
             // TODO: subscripts, etc.
             _ => {
                 self.infer_type_expression(slice);