diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs index 718093f51dbf6..ea0a931f851b2 100644 --- a/crates/red_knot_python_semantic/src/types.rs +++ b/crates/red_knot_python_semantic/src/types.rs @@ -157,12 +157,15 @@ impl<'db> Type<'db> { // TODO MRO? get_own_instance_member, get_instance_member todo!("attribute lookup on Instance type") } - Type::Union(_) => { - // TODO perform the get_member on each type in the union - // TODO return the union of those results - // TODO if any of those results is `None` then include Unknown in the result union - todo!("attribute lookup on Union type") - } + Type::Union(union) => Type::Union( + union + .elements(db) + .iter() + .fold(UnionTypeBuilder::new(db), |builder, element_ty| { + builder.add(element_ty.member(db, name)) + }) + .build(), + ), Type::Intersection(_) => { // TODO perform the get_member on each type in the intersection // TODO return the intersection of those results diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 355bba567313f..f8cd746401e3e 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -2238,6 +2238,28 @@ mod tests { Ok(()) } + #[test] + fn attribute_of_union() -> anyhow::Result<()> { + let mut db = setup_db(); + + db.write_dedented( + "/src/a.py", + " + if flag: + class C: + x = 1 + else: + class C: + x = 2 + y = C.x + ", + )?; + + assert_public_ty(&db, "/src/a.py", "y", "Literal[1, 2]"); + + Ok(()) + } + fn first_public_def<'db>(db: &'db TestDb, file: File, name: &str) -> Definition<'db> { let scope = global_scope(db, file); *use_def_map(db, scope)