[red-knot] implement attribute of union

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

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)