Convert union of types to type of union in unions_internal (#2708)

Summary:
Pull Request resolved: #2708

The spec says that `type` distributes over unions: https://typing.python.org/en/latest/spec/special-types.html#type, but we weren't always respecting this. Fixed by standardizing how we represent a union of types/type of union. I also fixed up a few places where we weren't handling `type[Union[...]]` correctly, which is a pre-existing issue that is more obvious now that we're standardizing to that form.

Differential Revision: D95638889

Author: rchen152

crates/pyrefly_types/src/simplify.rs

@@ -80,6 +80,7 @@ fn unions_internal(
             promote_anonymous_typed_dicts(&mut res, stdlib, heap);
         }
         collapse_tuple_unions_with_empty(&mut res, heap);
+        collapse_builtins_type(&mut res, heap);
         // `res` is collapsible again if `flatten_and_dedup` drops `xs` to 0 or 1 elements
         try_collapse(res, heap).unwrap_or_else(|members| heap.mk_union(members))
     })
@@ -320,6 +321,38 @@ fn flatten_unpacked_concrete_tuples(elts: Vec<Type>) -> Vec<Type> {
     result
 }
 
+/// `type[int] | type[str]` => `type[int | str]`
+fn collapse_builtins_type(types: &mut Vec<Type>, heap: &TypeHeap) {
+    let mut idx = 0;
+    let mut first_elt = None;
+    let mut additional_elts = Vec::new();
+    types.retain(|t| {
+        let retain = match t {
+            Type::Type(box t) if first_elt.is_none() => {
+                first_elt = Some((idx, t.clone()));
+                true
+            }
+            Type::Type(box t) => {
+                additional_elts.push(t.clone());
+                false
+            }
+            _ => true,
+        };
+        idx += 1;
+        retain
+    });
+    if let Some((idx, first_elt)) = first_elt
+        && !additional_elts.is_empty()
+    {
+        let mut elts = vec![first_elt.clone()];
+        elts.extend(additional_elts);
+        *(types
+            .get_mut(idx)
+            .expect("idx out of bounds when collapsing type members in union")) =
+            heap.mk_type_form(heap.mk_union(elts));
+    }
+}
+
 // After a TypeVarTuple gets substituted with a tuple type, try to simplify the type
 pub fn simplify_tuples(tuple: Tuple, _heap: &TypeHeap) -> Tuple {
     match tuple {

pyrefly/lib/alt/call.rs

@@ -234,6 +234,14 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             Type::Type(box Type::ClassType(cls)) => CallTargetLookup::Ok(Box::new(
                 CallTarget::Class(cls, ConstructorKind::TypeOfClass, None),
             )),
+            // `type[A | B]` is equivalent to `type[A] | type[B]` for call target resolution.
+            // Distribute `type[...]` over union members and resolve as a union.
+            Type::Type(box Type::Union(box Union { members: xs, .. })) => {
+                let union_of_types = self
+                    .heap
+                    .mk_union(xs.into_iter().map(|x| self.heap.mk_type_form(x)).collect());
+                self.as_call_target_impl(union_of_types, quantified)
+            }
             Type::Type(box Type::SelfType(cls)) => CallTargetLookup::Ok(Box::new(
                 CallTarget::Class(cls, ConstructorKind::TypeOfSelf, None),
             )),
@@ -326,7 +334,12 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             }
             Type::Any(style) => CallTargetLookup::Ok(Box::new(CallTarget::Any(style))),
             Type::TypeAlias(ta) => {
-                self.as_call_target_impl(self.get_type_alias(&ta).as_value(self.stdlib), quantified)
+                let body = self.get_type_alias(&ta).as_value(self.stdlib);
+                match body {
+                    // This comes from an expression like `int | str`, which is not callable.
+                    Type::Type(box Type::Union(_)) => CallTargetLookup::Error(vec![]),
+                    _ => self.as_call_target_impl(body, quantified),
+                }
             }
             Type::ClassType(cls) => {
                 let maybe_dunder_call = if let Some(quantified) = &quantified {

pyrefly/lib/report/pysa/types.rs

@@ -292,6 +292,20 @@ fn get_classes_of_type(type_: &Type, context: &ModuleContext) -> ClassNamesFromT
             ClassNamesFromType::from_class(class_type.class_object(), context)
                 .prepend_modifier(TypeModifier::Type)
         }
+        Type::Type(box Type::Union(box Union {
+            members: elements, ..
+        })) if !elements.is_empty() => elements
+            .iter()
+            .map(|inner| match inner {
+                Type::ClassType(class_type) => {
+                    ClassNamesFromType::from_class(class_type.class_object(), context)
+                        .prepend_modifier(TypeModifier::Type)
+                }
+                _ => ClassNamesFromType::not_a_class(),
+            })
+            .reduce(|acc, next| acc.join_with(next))
+            .expect("expected at least one element in union")
+            .sort_and_dedup(),
         Type::Tuple(_) => ClassNamesFromType::from_class(context.stdlib.tuple_object(), context),
         Type::TypedDict(TypedDict::TypedDict(inner)) => {
             ClassNamesFromType::from_class(inner.class_object(), context)

pyrefly/lib/test/pysa/types.rs

@@ -492,7 +492,7 @@ class MyTypedDict(TypedDict):
 
     assert_eq!(
         PysaType::new(
-            "type[test.A] | type[test.B]".to_owned(),
+            "type[test.A | test.B]".to_owned(),
             ClassNamesFromType::from_classes(
                 vec![
                     get_class_ref("test", "A", &context),

pyrefly/lib/test/simple.rs

@@ -58,6 +58,54 @@ def f(x: type[int] | type[str], y: type[int | str]) -> None:
 "#,
 );
 
+testcase!(
+    test_distribute_type_assignability,
+    r#"
+def f() -> type[int | str]: ...
+def g() -> type[int] | type[str]: ...
+x: type[int] | type[str] = f()
+y: type[int | str] = g()
+    "#,
+);
+
+testcase!(
+    test_type_of_union_matches_type_param,
+    r#"
+from typing import assert_type
+def f[T](x: type[T]) -> T: ...
+assert_type(f(int | None), int | None)
+    "#,
+);
+
+testcase!(
+    test_type_of_union_matches_type_param_or_none,
+    r#"
+from typing import assert_type
+def f[T](x: type[T] | None) -> T: ...
+assert_type(f(int | str), int | str)
+    "#,
+);
+
+testcase!(
+    test_type_of_union_partially_matches_type_param,
+    r#"
+from typing import assert_type
+def f[T](x: type[T] | type[int]) -> T: ...
+def g(x: type[int | str]):
+    assert_type(f(x), str)
+    "#,
+);
+
+testcase!(
+    test_union_of_type_matches_type_param,
+    r#"
+from typing import assert_type
+def f[T](x: type[T]) -> T: ...
+def g(x: type[int] | type[str]):
+    assert_type(f(x), int | str)
+    "#,
+);
+
 testcase!(
     test_simple_call,
     r#"

pyrefly/lib/test/type_alias.rs

@@ -1262,3 +1262,12 @@ def f3(x: C) -> None:
     assert_type(x, list[int])
     "#,
 );
+
+testcase!(
+    test_union_is_not_callable,
+    r#"
+from typing import TypeAlias
+X: TypeAlias = int | str
+X()  # E: Expected a callable
+    "#,
+);