Skip to content

Commit ee0518e

Browse files
authored
[red-knot] implement attribute of union (#12601)
I hit this `todo!` trying to run type inference over some real modules. Since it's a one-liner to implement it, I just did that rather than changing to `Type::Unknown`.
1 parent d774a3b commit ee0518e

File tree

2 files changed

+31
-6
lines changed

2 files changed

+31
-6
lines changed

crates/red_knot_python_semantic/src/types.rs

+9-6
Original file line numberDiff line numberDiff line change
@@ -157,12 +157,15 @@ impl<'db> Type<'db> {
157157
// TODO MRO? get_own_instance_member, get_instance_member
158158
todo!("attribute lookup on Instance type")
159159
}
160-
Type::Union(_) => {
161-
// TODO perform the get_member on each type in the union
162-
// TODO return the union of those results
163-
// TODO if any of those results is `None` then include Unknown in the result union
164-
todo!("attribute lookup on Union type")
165-
}
160+
Type::Union(union) => Type::Union(
161+
union
162+
.elements(db)
163+
.iter()
164+
.fold(UnionTypeBuilder::new(db), |builder, element_ty| {
165+
builder.add(element_ty.member(db, name))
166+
})
167+
.build(),
168+
),
166169
Type::Intersection(_) => {
167170
// TODO perform the get_member on each type in the intersection
168171
// TODO return the intersection of those results

crates/red_knot_python_semantic/src/types/infer.rs

+22
Original file line numberDiff line numberDiff line change
@@ -2238,6 +2238,28 @@ mod tests {
22382238
Ok(())
22392239
}
22402240

2241+
#[test]
2242+
fn attribute_of_union() -> anyhow::Result<()> {
2243+
let mut db = setup_db();
2244+
2245+
db.write_dedented(
2246+
"/src/a.py",
2247+
"
2248+
if flag:
2249+
class C:
2250+
x = 1
2251+
else:
2252+
class C:
2253+
x = 2
2254+
y = C.x
2255+
",
2256+
)?;
2257+
2258+
assert_public_ty(&db, "/src/a.py", "y", "Literal[1, 2]");
2259+
2260+
Ok(())
2261+
}
2262+
22412263
fn first_public_def<'db>(db: &'db TestDb, file: File, name: &str) -> Definition<'db> {
22422264
let scope = global_scope(db, file);
22432265
*use_def_map(db, scope)

0 commit comments

Comments
 (0)