[red-knot] Fix: Infer type for typing.Union[..] tuple expression (#14510

) ## Summary Fixes a panic related to sub-expressions of `typing.Union` where we fail to store a type for the `int, str` tuple-expression in code like this: ``` x: Union[int, str] = 1 ``` relates to [my comment](#14499 (comment)) on #14499. ## Test Plan New corpus test
astral-sh · Nov 21, 2024 · f684b6f · f684b6f
1 parent 47f39ed
commit f684b6f
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -4568,10 +4568,14 @@ impl<'db> TypeInferenceBuilder<'db> {
                 UnionType::from_elements(self.db, [param_type, Type::none(self.db)])
             }
             KnownInstanceType::Union => match parameters {
-                ast::Expr::Tuple(t) => UnionType::from_elements(
-                    self.db,
-                    t.iter().map(|elt| self.infer_type_expression(elt)),
-                ),
+                ast::Expr::Tuple(t) => {
+                    let union_ty = UnionType::from_elements(
+                        self.db,
+                        t.iter().map(|elt| self.infer_type_expression(elt)),
+                    );
+                    self.store_expression_type(parameters, union_ty);
+                    union_ty
+                }
                 _ => self.infer_type_expression(parameters),
             },
             KnownInstanceType::TypeVar(_) => todo_type!(),

diff --git a/crates/red_knot_workspace/resources/test/corpus/95_annotation_union.py b/crates/red_knot_workspace/resources/test/corpus/95_annotation_union.py
@@ -0,0 +1,3 @@
+from typing import Union
+
+x: Union[int, str] = 1
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		from typing import Union

		x: Union[int, str] = 1