From ee0518e8f71cb0bc4d5cc3afc28a54a0bbc78247 Mon Sep 17 00:00:00 2001
From: Carl Meyer <carl@astral.sh>
Date: Wed, 31 Jul 2024 19:45:24 -0700
Subject: [PATCH] [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`.
---
 crates/red_knot_python_semantic/src/types.rs  | 15 ++++++++-----
 .../src/types/infer.rs                        | 22 +++++++++++++++++++
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs
index 718093f51d..ea0a931f85 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 355bba5673..f8cd746401 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)