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)

Reviewed By: rchen152

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

@@ -449,11 +449,8 @@ impl<'a> TypeDisplayContext<'a> {
             Type::Union(box Union {
                 display_name: Some(name),
                 ..
-            }) => output.write_str(name),
-            Type::Union(box Union {
-                members,
-                display_name: None,
-            }) => {
+            }) if !is_toplevel => output.write_str(name),
+            Type::Union(box Union { members, .. }) => {
                 let mut literal_idx = None;
                 let mut literals = Vec::new();
                 let mut union_members: Vec<&Type> = Vec::new();

crates/pyrefly_types/src/types.rs

@@ -1472,10 +1472,15 @@ 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 Union { members: ts, .. }) = ty {
+            if let Type::Union(box Union {
+                members: ts,
+                display_name,
+            }) = ty
+            {
                 ts.sort();
+                *display_name = None;
             }
         })
     }

pyrefly/lib/alt/solve.rs

@@ -3800,7 +3800,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 Union { display_name, .. }) = &mut aliased_type {
+                    *display_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

@@ -70,8 +70,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

@@ -457,8 +457,17 @@ 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 Union { members: xs, .. }) = x {
-                *x = unions(mem::take(xs));
+            if let Type::Union(box Union {
+                members: xs,
+                display_name: original_name,
+            }) = x
+            {
+                let mut merged = unions(mem::take(xs));
+                // Preserve union display names during simplification
+                if let Type::Union(box Union { display_name, .. }) = &mut merged {
+                    *display_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

@@ -824,3 +824,53 @@ class UseC(Generic[T]):
             current.make()
     "#,
 );
+
+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`
+    "#,
+);