add optional display name to Type::Union

Summary:
This is an alternative to D84849556, which aims to solve facebook#1219

Instead of introducing a new type variant, this diff introduces a new field in Type::Union to hold an optional display name. It will be populated when resolving a type alias.

Compared to the other approach:

benefits
- less plumbing
- type checking behavior unaffected

drawbacks
- only applies to type aliases of unions (though arguably this is where readability matters most)
- name information is not preserved during union flattening (so if we union a type alias with something else, the result is the full expanded type)

Differential Revision: D87558541

Reviewed By: rchen152

Author: yangdanny97

crates/pyrefly_types/src/display.rs

@@ -40,6 +40,7 @@ use crate::types::SuperObj;
 use crate::types::TArgs;
 use crate::types::TParam;
 use crate::types::Type;
+use crate::types::Union;
 
 /// Information about the qnames we have seen.
 /// Set to None to indicate we have seen different values, or Some if they are all the same.
@@ -442,17 +443,24 @@ impl<'a> TypeDisplayContext<'a> {
                 self.maybe_fmt_with_module("typing", "NoReturn", output)
             }
             Type::Never(NeverStyle::Never) => self.maybe_fmt_with_module("typing", "Never", output),
-            Type::Union(types) if types.is_empty() => {
+            Type::Union(box Union { members: types, .. }) if types.is_empty() => {
                 self.maybe_fmt_with_module("typing", "Never", output)
             }
-            Type::Union(types) => {
+            Type::Union(box Union {
+                display_name: Some(name),
+                ..
+            }) => output.write_str(name),
+            Type::Union(box Union {
+                members,
+                display_name: None,
+            }) => {
                 let mut literal_idx = None;
                 let mut literals = Vec::new();
                 let mut union_members: Vec<&Type> = Vec::new();
                 // Track seen types to deduplicate (mainly to prettify types for functions with different names but the same signature)
                 let mut seen_types = SmallSet::new();
 
-                for t in types.iter() {
+                for t in members.iter() {
                     match t {
                         Type::Literal(lit) => {
                             if literal_idx.is_none() {
@@ -1099,7 +1107,7 @@ pub mod tests {
             assert_eq!(ctx.display(&int_type).to_string(), "int");
         }
 
-        let union_foo_int = Type::Union(vec![foo_type, int_type]);
+        let union_foo_int = Type::union(vec![foo_type, int_type]);
 
         {
             let mut ctx = TypeDisplayContext::new(&[&union_foo_int]);
@@ -1115,11 +1123,11 @@ pub mod tests {
         let t3 = fake_tyvar("qux", "bar", 2);
 
         assert_eq!(
-            Type::Union(vec![t1.to_type(), t2.to_type()]).to_string(),
+            Type::union(vec![t1.to_type(), t2.to_type()]).to_string(),
             "TypeVar[bar.foo@1:2] | TypeVar[bar.foo@1:3]"
         );
         assert_eq!(
-            Type::Union(vec![t1.to_type(), t3.to_type()]).to_string(),
+            Type::union(vec![t1.to_type(), t3.to_type()]).to_string(),
             "TypeVar[foo] | TypeVar[qux]"
         );
     }
@@ -1159,13 +1167,14 @@ pub mod tests {
         let nonlit2 = Type::LiteralString;
 
         assert_eq!(
-            Type::Union(vec![nonlit1.clone(), nonlit2.clone()]).to_string(),
+            Type::union(vec![nonlit1.clone(), nonlit2.clone()]).to_string(),
             "None | LiteralString"
         );
         assert_eq!(
-            Type::Union(vec![nonlit1, lit1, nonlit2, lit2]).to_string(),
+            Type::union(vec![nonlit1.clone(), lit1, nonlit2.clone(), lit2]).to_string(),
             "None | Literal[True, 'test'] | LiteralString"
         );
+        assert_eq!(Type::union(vec![nonlit1, nonlit2]).to_string(), "MyUnion");
     }
 
     #[test]
@@ -1575,7 +1584,7 @@ def overloaded_func[T](
 
     #[test]
     fn test_union_of_intersection() {
-        let x = Type::Union(vec![
+        let x = Type::union(vec![
             Type::Intersect(Box::new((
                 vec![Type::any_explicit(), Type::LiteralString],
                 Type::any_implicit(),

crates/pyrefly_types/src/simplify.rs

@@ -14,13 +14,14 @@ use crate::literal::Lit;
 use crate::stdlib::Stdlib;
 use crate::tuple::Tuple;
 use crate::types::Type;
+use crate::types::Union;
 
 /// Turn unions of unions into a flattened list for one union, and return the deduped list.
 fn flatten_and_dedup(xs: Vec<Type>) -> Vec<Type> {
     fn flatten(xs: Vec<Type>, res: &mut Vec<Type>) {
         for x in xs {
             match x {
-                Type::Union(xs) => flatten(xs, res),
+                Type::Union(box Union { members, .. }) => flatten(members, res),
                 Type::Never(_) => {}
                 _ => res.push(x),
             }
@@ -76,7 +77,12 @@ fn unions_internal(
         }
         collapse_tuple_unions_with_empty(&mut res);
         // `res` is collapsible again if `flatten_and_dedup` drops `xs` to 0 or 1 elements
-        try_collapse(res).unwrap_or_else(Type::Union)
+        try_collapse(res).unwrap_or_else(|members| {
+            Type::Union(Box::new(Union {
+                members,
+                display_name: None,
+            }))
+        })
     })
 }
 

crates/pyrefly_types/src/type_output.rs

@@ -402,7 +402,7 @@ mod tests {
 
         let int_type = Type::ClassType(ClassType::new(int_class, TArgs::default()));
         let str_type = Type::ClassType(ClassType::new(str_class, TArgs::default()));
-        let union_type = Type::Union(vec![int_type, str_type, Type::None]);
+        let union_type = Type::union(vec![int_type, str_type, Type::None]);
 
         let ctx = TypeDisplayContext::new(&[&union_type]);
         let mut output = OutputWithLocations::new(&ctx);

crates/pyrefly_types/src/types.rs

@@ -9,6 +9,8 @@ use std::borrow::Cow;
 use std::cmp::Ordering;
 use std::fmt;
 use std::fmt::Display;
+use std::hash::Hash;
+use std::hash::Hasher;
 use std::sync::Arc;
 
 use dupe::Dupe;
@@ -611,6 +613,40 @@ pub enum SuperObj {
     Class(ClassType),
 }
 
+#[derive(Debug, Clone, Eq, TypeEq, PartialOrd, Ord)]
+pub struct Union {
+    pub members: Vec<Type>,
+    pub display_name: Option<String>,
+}
+
+impl PartialEq for Union {
+    fn eq(&self, other: &Self) -> bool {
+        self.members == other.members
+    }
+}
+
+impl Hash for Union {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.members.hash(state)
+    }
+}
+
+impl Visit<Type> for Union {
+    fn recurse<'a>(&'a self, f: &mut dyn FnMut(&'a Type)) {
+        for member in &self.members {
+            member.visit(f);
+        }
+    }
+}
+
+impl VisitMut<Type> for Union {
+    fn recurse_mut(&mut self, f: &mut dyn FnMut(&mut Type)) {
+        for member in &mut self.members {
+            member.visit_mut(f);
+        }
+    }
+}
+
 // Note: The fact that Literal and LiteralString are at the front is important for
 // optimisations in `unions_with_literals`.
 #[derive(Debug, Clone, PartialEq, Eq, TypeEq, PartialOrd, Ord, Hash)]
@@ -626,7 +662,8 @@ pub enum Type {
     BoundMethod(Box<BoundMethod>),
     /// An overloaded function.
     Overload(Overload),
-    Union(Vec<Type>),
+    /// Unions will hold an optional name to use when displaying the type
+    Union(Box<Union>),
     /// Our intersection support is partial, so we store a fallback type that we use for operations
     /// that are not yet supported on intersections.
     Intersect(Box<(Vec<Type>, Type)>),
@@ -1392,7 +1429,7 @@ impl Type {
     /// type[a | b] -> type[a] | type[b]
     pub fn distribute_type_over_union(self) -> Self {
         self.transform(&mut |ty| {
-            if let Type::Type(box Type::Union(members)) = ty {
+            if let Type::Type(box Type::Union(box Union { members, .. })) = ty {
                 *ty = unions(members.drain(..).map(Type::type_form).collect());
             }
         })
@@ -1437,7 +1474,7 @@ impl Type {
 
     pub fn sort_unions(self) -> Self {
         self.transform(&mut |ty| {
-            if let Type::Union(ts) = ty {
+            if let Type::Union(box Union { members: ts, .. }) = ty {
                 ts.sort();
             }
         })
@@ -1488,27 +1525,30 @@ impl Type {
 
     pub fn into_unions(self) -> Vec<Type> {
         match self {
-            Type::Union(types) => types,
+            Type::Union(box Union { members: types, .. }) => types,
             _ => vec![self],
         }
     }
 
     /// Create an optional type (union with None).
     pub fn optional(x: Self) -> Self {
         // We would like the resulting type not nested, and well sorted.
-        if let Type::Union(mut xs) = x {
+        if let Type::Union(box Union {
+            members: mut xs, ..
+        }) = x
+        {
             match xs.binary_search(&Type::None) {
-                Ok(_) => Type::Union(xs),
+                Ok(_) => Type::union(xs),
                 Err(i) => {
                     xs.insert(i, Type::None);
-                    Type::Union(xs)
+                    Type::union(xs)
                 }
             }
         } else {
             match x.cmp(&Type::None) {
                 Ordering::Equal => Type::None,
-                Ordering::Less => Type::Union(vec![x, Type::None]),
-                Ordering::Greater => Type::Union(vec![Type::None, x]),
+                Ordering::Less => Type::union(vec![x, Type::None]),
+                Ordering::Greater => Type::union(vec![Type::None, x]),
             }
         }
     }
@@ -1538,9 +1578,9 @@ impl Type {
             Type::Literal(Lit::Str(x)) => Some(!x.is_empty()),
             Type::None => Some(false),
             Type::Tuple(Tuple::Concrete(elements)) => Some(!elements.is_empty()),
-            Type::Union(options) => {
+            Type::Union(box Union { members, .. }) => {
                 let mut answer = None;
-                for option in options {
+                for option in members {
                     let option_bool = option.as_bool();
                     option_bool?;
                     if answer.is_none() {
@@ -1590,6 +1630,14 @@ impl Type {
             })
         })
     }
+
+    /// Creates a union from the provided types without simplifying
+    pub fn union(members: Vec<Type>) -> Self {
+        Type::Union(Box::new(Union {
+            members,
+            display_name: None,
+        }))
+    }
 }
 
 #[cfg(test)]
@@ -1616,8 +1664,8 @@ mod tests {
         let false_lit = Type::Literal(Lit::Bool(false));
         let none = Type::None;
 
-        let str_opt = Type::Union(vec![s, none.clone()]);
-        let false_opt = Type::Union(vec![false_lit, none]);
+        let str_opt = Type::union(vec![s, none.clone()]);
+        let false_opt = Type::union(vec![false_lit, none]);
 
         assert_eq!(str_opt.as_bool(), None);
         assert_eq!(false_opt.as_bool(), Some(false));

pyrefly/lib/alt/answers_solver.rs

@@ -17,6 +17,7 @@ use dupe::IterDupedExt;
 use itertools::Either;
 use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::module_path::ModulePath;
+use pyrefly_types::types::Union;
 use pyrefly_util::display::DisplayWithCtx;
 use pyrefly_util::display::commas_iter;
 use pyrefly_util::recurser::Guard;
@@ -797,12 +798,12 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             fn go(&mut self, ty: &Type, in_type: bool) {
                 match ty {
                     Type::Never(_) if !in_type => (),
-                    Type::Union(tys) => {
+                    Type::Union(box Union { members, .. }) => {
                         self.seen_union = true;
-                        tys.iter().for_each(|ty| self.go(ty, in_type))
+                        members.iter().for_each(|ty| self.go(ty, in_type))
                     }
-                    Type::Type(box Type::Union(tys)) if !in_type => {
-                        tys.iter().for_each(|ty| self.go(ty, true))
+                    Type::Type(box Type::Union(box Union { members, .. })) if !in_type => {
+                        members.iter().for_each(|ty| self.go(ty, true))
                     }
                     Type::Var(v) if let Some(_guard) = self.me.recurse(*v) => {
                         self.go(&self.me.solver().force_var(*v), in_type)

pyrefly/lib/alt/attr.rs

@@ -16,6 +16,7 @@ use pyrefly_types::special_form::SpecialForm;
 use pyrefly_types::types::Forall;
 use pyrefly_types::types::Forallable;
 use pyrefly_types::types::TArgs;
+use pyrefly_types::types::Union;
 use pyrefly_types::types::Var;
 use ruff_python_ast::helpers::is_dunder;
 use ruff_python_ast::name::Name;
@@ -1795,12 +1796,12 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             Type::SuperInstance(box (cls, obj)) => {
                 acc.push(AttributeBase1::SuperInstance(cls, obj))
             }
-            Type::Union(members) => {
+            Type::Union(box Union { members, .. }) => {
                 for ty in members {
                     self.as_attribute_base1(ty, acc)
                 }
             }
-            Type::Type(box Type::Union(members)) => {
+            Type::Type(box Type::Union(box Union { members, .. })) => {
                 for ty in members {
                     self.as_attribute_base1(Type::type_form(ty), acc)
                 }

pyrefly/lib/alt/call.rs

@@ -14,6 +14,7 @@ use pyrefly_types::quantified::Quantified;
 use pyrefly_types::types::CalleeKind;
 use pyrefly_types::types::TArgs;
 use pyrefly_types::types::TParams;
+use pyrefly_types::types::Union;
 use pyrefly_util::prelude::SliceExt;
 use pyrefly_util::prelude::VecExt;
 use ruff_python_ast::Arguments;
@@ -260,7 +261,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             Type::Var(v) if let Some(_guard) = self.recurse(v) => {
                 self.as_call_target_impl(self.solver().force_var(v), quantified, dunder_call)
             }
-            Type::Union(xs) => {
+            Type::Union(box Union { members: xs, .. }) => {
                 let xs_length = xs.len();
                 let targets = xs
                     .into_iter()
@@ -315,7 +316,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             Type::Quantified(q) if q.is_type_var() => match q.restriction() {
                 Restriction::Unrestricted => CallTargetLookup::Error(vec![]),
                 Restriction::Bound(bound) => match bound {
-                    Type::Union(members) => {
+                    Type::Union(box Union { members, .. }) => {
                         let mut targets = Vec::new();
                         for member in members {
                             if let CallTargetLookup::Ok(target) = self.as_call_target_impl(

pyrefly/lib/alt/class/class_field.rs

@@ -22,6 +22,7 @@ use pyrefly_types::callable::Params;
 use pyrefly_types::simplify::unions;
 use pyrefly_types::type_var::Restriction;
 use pyrefly_types::types::TParams;
+use pyrefly_types::types::Union;
 use pyrefly_util::owner::Owner;
 use pyrefly_util::prelude::ResultExt;
 use pyrefly_util::visit::Visit;
@@ -944,7 +945,9 @@ fn make_bound_method_helper(
         Type::Overload(overload) if should_bind2(&overload.metadata) => {
             BoundMethodType::Overload(overload)
         }
-        Type::Union(ref ts) => {
+        Type::Union(box Union {
+            members: ref ts, ..
+        }) => {
             let mut bound_methods = Vec::with_capacity(ts.len());
             for t in ts {
                 match make_bound_method_helper(obj.clone(), t.clone(), should_bind) {

pyrefly/lib/alt/class/classdef.rs

@@ -11,6 +11,7 @@ use dupe::Dupe;
 use pyrefly_python::nesting_context::NestingContext;
 use pyrefly_types::callable::Callable;
 use pyrefly_types::special_form::SpecialForm;
+use pyrefly_types::types::Union;
 use ruff_python_ast::Identifier;
 use ruff_python_ast::StmtClassDef;
 use ruff_python_ast::name::Name;
@@ -51,7 +52,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             Type::SelfType(ty_cls) | Type::ClassType(ty_cls) => {
                 self.has_superclass(ty_cls.class_object(), class)
             }
-            Type::Union(xs) => xs
+            Type::Union(box Union { members: xs, .. }) => xs
                 .iter()
                 .all(|x| self.is_compatible_constructor_return(x, class)),
             _ => false,