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

sharkdp · web-flow · commit f684b6fff4c1 · 2024-11-21T11:49:20.000+01:00
## 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
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 @@`
	`1`	`+from typing import Union`
	`2`	`+`
	`3`	`+x: Union[int, str] = 1`