populate display name when untyping a type alias (facebook#1641)

Summary:
Pull Request resolved: facebook#1641

- populate display name when untyping a union
- keep name when simplifying
- display the expanded type when the top level type being printed is the alias (only use the display name for nested types)

Differential Revision: D87558543

Author: yangdanny97

conformance/third_party/conformance.exp

@@ -15,8 +15,8 @@
     {
       "code": -2,
       "column": 9,
-      "concise_description": "`TypeAlias[GoodTypeAlias3, type[list[int | None]]]` is not subscriptable",
-      "description": "`TypeAlias[GoodTypeAlias3, type[list[int | None]]]` is not subscriptable",
+      "concise_description": "`TypeAlias[GoodTypeAlias3, type[list[GoodTypeAlias2]]]` is not subscriptable",
+      "description": "`TypeAlias[GoodTypeAlias3, type[list[GoodTypeAlias2]]]` is not subscriptable",
       "line": 68,
       "name": "unsupported-operation",
       "severity": "error",
@@ -270,8 +270,8 @@
     {
       "code": -2,
       "column": 9,
-      "concise_description": "`type[list[int | None]]` is not subscriptable",
-      "description": "`type[list[int | None]]` is not subscriptable",
+      "concise_description": "`type[list[GoodTypeAlias2]]` is not subscriptable",
+      "description": "`type[list[GoodTypeAlias2]]` is not subscriptable",
       "line": 77,
       "name": "unsupported-operation",
       "severity": "error",

crates/pyrefly_types/src/display.rs

@@ -415,8 +415,8 @@ impl<'a> TypeDisplayContext<'a> {
             Type::Union(box (types, _)) if types.is_empty() => {
                 self.maybe_fmt_with_module("typing", "Never", output)
             }
-            Type::Union(box (_, Some(name))) => output.write_str(name),
-            Type::Union(box (types, None)) => {
+            Type::Union(box (_, Some(name))) if !is_toplevel => output.write_str(name),
+            Type::Union(box (types, _)) => {
                 // All Literals will be collected into a single Literal at the index of the first Literal.
                 let mut literal_idx = None;
                 let mut literals = Vec::new();
@@ -1100,13 +1100,24 @@ pub mod tests {
             Type::union(vec![nonlit1.clone(), lit1, nonlit2.clone(), lit2]).to_string(),
             "None | Literal[True, 'test'] | LiteralString"
         );
+
+        // Unions have an optional name field, which is shown only when
+        // the union is nested in another type
         assert_eq!(
             Type::Union(Box::new((
-                vec![nonlit1, nonlit2],
+                vec![nonlit1.clone(), nonlit2.clone()],
                 Some("MyUnion".to_owned())
             )))
             .to_string(),
-            "MyUnion"
+            "None | LiteralString"
+        );
+        assert_eq!(
+            Type::Type(Box::new(Type::Union(Box::new((
+                vec![nonlit1, nonlit2],
+                Some("MyUnion".to_owned())
+            )))))
+            .to_string(),
+            "type[MyUnion]"
         );
     }
 

crates/pyrefly_types/src/types.rs

@@ -1422,10 +1422,11 @@ impl Type {
         })
     }
 
-    pub fn sort_unions(self) -> Self {
+    pub fn sort_unions_and_drop_names(self) -> Self {
         self.transform(&mut |ty| {
-            if let Type::Union(box (ts, _)) = ty {
+            if let Type::Union(box (ts, name)) = ty {
                 ts.sort();
+                *name = None;
             }
         })
     }

pyrefly/lib/alt/solve.rs

@@ -3751,7 +3751,13 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             Type::None => Some(Type::None), // Both a value and a type
             Type::Ellipsis => Some(Type::Ellipsis), // A bit weird because of tuples, so just promote it
             Type::Any(style) => Some(style.propagate()),
-            Type::TypeAlias(ta) => self.untype_opt(ta.as_type(), range, errors),
+            Type::TypeAlias(ta) => {
+                let mut aliased_type = self.untype_opt(ta.as_type(), range, errors)?;
+                if let Type::Union(box (_, name)) = &mut aliased_type {
+                    *name = Some(ta.name.to_string());
+                }
+                Some(aliased_type)
+            }
             t @ Type::Unpack(
                 box Type::Tuple(_) | box Type::TypeVarTuple(_) | box Type::Quantified(_),
             ) => Some(t),

pyrefly/lib/alt/special_calls.rs

@@ -69,8 +69,9 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                     .distribute_type_over_union();
                 // Make assert_type(Self@SomeClass, typing.Self) work.
                 ty.subst_self_type_mut(&self_form);
-                // Re-sort unions. Make sure to keep this as the final step before comparison.
-                ty.sort_unions()
+                // Re-sort unions & drop any display names.
+                // Make sure to keep this as the final step before comparison.
+                ty.sort_unions_and_drop_names()
             };
             let a = normalize_type(a, expr_a);
             let b = normalize_type(b, expr_b);

pyrefly/lib/solver/solver.rs

@@ -456,8 +456,13 @@ impl Solver {
     /// Simplify a type as much as we can.
     fn simplify_mut(&self, t: &mut Type) {
         t.transform_mut(&mut |x| {
-            if let Type::Union(box (xs, _)) = x {
-                *x = unions(mem::take(xs));
+            if let Type::Union(box (xs, original_name)) = x {
+                let mut merged = unions(mem::take(xs));
+                // Preserve union display names during simplification
+                if let Type::Union(box (_, name)) = &mut merged {
+                    *name = original_name.clone();
+                }
+                *x = merged;
             }
             if let Type::Intersect(y) = x {
                 *x = intersect(mem::take(&mut y.0), y.1.clone());

pyrefly/lib/test/class_overrides.rs

@@ -577,7 +577,7 @@ class A:
     def f(self, x: TA1):
         pass
 class B(A):
-    def f(self, x: TA2):  # E: `B.f` has type `BoundMethod[B, (self: B, x: float | int) -> None]`, which is not assignable to `BoundMethod[B, (self: B, x: float | int | None) -> None]`, the type of `A.f`
+    def f(self, x: TA2):  # E: `B.f` has type `BoundMethod[B, (self: B, x: TA2) -> None]`, which is not assignable to `BoundMethod[B, (self: B, x: TA1) -> None]`, the type of `A.f`
         pass
     "#,
 );

pyrefly/lib/test/type_alias.rs

@@ -806,3 +806,53 @@ x1: Spam1[int, str] = int
 x2: Spam2[int, str] = int
     "#,
 );
+
+fn env_with_alias() -> TestEnv {
+    TestEnv::one(
+        "foo",
+        r#"
+type TA = int | str
+
+def f(x: TA) -> TA:
+  return x
+"#,
+    )
+}
+
+testcase!(
+    test_alias_union_name,
+    env_with_alias(),
+    r#"
+from foo import TA, f
+from typing import Callable
+
+val1: int | str = 1
+val2: TA = 1
+
+# Union names are only shown when nested in another type
+
+f(object())  # E: Argument `object` is not assignable to parameter `x` with type `int | str` in function `foo.f`
+f(val1)
+f(val2)
+x1: TA = object()  # E: `object` is not assignable to `int | str`
+x2: TA = val1
+x3: TA = val2
+c1: Callable[[int], int] = f  # E: `(x: TA) -> TA` is not assignable to `(int) -> int`
+
+# Union names are lost when flattened into another union
+
+class C: pass
+def f2(x: TA | C) -> TA | C:
+  return x
+
+f2(object())  # E: Argument `object` is not assignable to parameter `x` with type `C | int | str` in function `f2`
+f2(val1)
+f2(val2)
+f2(C())
+x4: TA | C = object()  # E: `object` is not assignable to `C | int | str`
+x5: TA | C = val1
+x6: TA | C = val2
+x7: TA | C = C()
+c2: Callable[[int], int] = f2  # E: `(x: C | int | str) -> C | int | str` is not assignable to `(int) -> int`
+    "#,
+);