From 269e47be961d3eb95c0351e8f7006b71e083a6a4 Mon Sep 17 00:00:00 2001
From: Shaygan Hooshyari <sh.hooshyari@gmail.com>
Date: Sat, 7 Dec 2024 18:34:50 +0100
Subject: [PATCH] Understand `type[A | B]` special form in annotations (#14830)

resolves https://github.com/astral-sh/ruff/issues/14703

I decided to use recursion to get the type, so if anything is added to
the single element inference it will be applied for the union.
Also added this
[change](https://github.com/astral-sh/ruff/issues/14703#issuecomment-2510286217)
in this PR since it was easy.

---------

Co-authored-by: Carl Meyer <carl@astral.sh>
---
 .../resources/mdtest/type_of/basic.md         | 28 +++++++++++++++++++
 .../src/types/infer.rs                        | 23 ++++++++++++++-
 2 files changed, 50 insertions(+), 1 deletion(-)

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")