feat(formatter): implement formatting for TSUnionType

crates/oxc_formatter/src/formatter/comments.rs

@@ -173,6 +173,27 @@ impl<'a> Comments<'a> {
         &self.unprinted_comments()[..index]
     }
 
+    /// Returns end-of-line comments that are after the given position (excluding printed ones).
+    pub fn end_of_line_comments_after(&self, mut pos: u32) -> &'a [Comment] {
+        let comments = self.unprinted_comments();
+        for (index, comment) in comments.iter().enumerate() {
+            if self
+                .source_text
+                .all_bytes_match(pos, comment.span.start, |b| matches!(b, b'\t' | b' ' | b')'))
+            {
+                if !self.source_text.is_own_line_comment(comment)
+                    && (comment.is_line() || self.source_text.is_end_of_line_comment(comment))
+                {
+                    return &comments[..=index];
+                }
+                pos = comment.span.end;
+            } else {
+                break;
+            }
+        }
+        &[]
+    }
+
     /// Returns comments that start after the given position (excluding printed ones).
     pub fn comments_after(&self, pos: u32) -> &'a [Comment] {
         let comments = self.unprinted_comments();

crates/oxc_formatter/src/generated/format.rs

@@ -3786,9 +3786,19 @@ impl<'a> Format<'a> for AstNode<'a, TSConditionalType<'a>> {
 impl<'a> Format<'a> for AstNode<'a, TSUnionType<'a>> {
     fn fmt(&self, f: &mut Formatter<'_, 'a>) -> FormatResult<()> {
         let is_suppressed = f.comments().is_suppressed(self.span().start);
+        if !is_suppressed && format_type_cast_comment_node(self, false, f)? {
+            return Ok(());
+        }
         self.format_leading_comments(f)?;
+        let needs_parentheses = self.needs_parentheses(f);
+        if needs_parentheses {
+            "(".fmt(f)?;
+        }
         let result =
             if is_suppressed { FormatSuppressedNode(self.span()).fmt(f) } else { self.write(f) };
+        if needs_parentheses {
+            ")".fmt(f)?;
+        }
         self.format_trailing_comments(f)?;
         result
     }

crates/oxc_formatter/src/parentheses/ts_type.rs

@@ -60,12 +60,22 @@ impl<'a> NeedsParentheses<'a> for AstNode<'a, TSConstructorType<'a>> {
 
 impl<'a> NeedsParentheses<'a> for AstNode<'a, TSUnionType<'a>> {
     fn needs_parentheses(&self, f: &Formatter<'_, 'a>) -> bool {
-        matches!(
-            self.parent,
-            AstNodes::TSArrayType(_)
-                | AstNodes::TSTypeOperator(_)
-                | AstNodes::TSIndexedAccessType(_)
-        )
+        match self.parent {
+            AstNodes::TSUnionType(union) => self.types.len() > 1 && union.types.len() > 1,
+            AstNodes::TSIntersectionType(intersection) => {
+                self.types.len() > 1 && intersection.types.len() > 1
+            }
+            parent => operator_type_or_higher_needs_parens(self.span(), parent),
+        }
+    }
+}
+
+/// Returns `true` if a TS primary type needs parentheses
+fn operator_type_or_higher_needs_parens(span: Span, parent: &AstNodes) -> bool {
+    match parent {
+        AstNodes::TSArrayType(_) | AstNodes::TSTypeOperator(_) | AstNodes::TSRestType(_) => true,
+        AstNodes::TSIndexedAccessType(indexed) => indexed.object_type.span() == span,
+        _ => false,
     }
 }
 

crates/oxc_formatter/src/write/mod.rs

@@ -26,6 +26,7 @@ mod switch_statement;
 mod template;
 mod try_statement;
 mod type_parameters;
+mod union_type;
 mod utils;
 mod variable_declaration;
 
@@ -1209,26 +1210,6 @@ impl<'a> FormatWrite<'a> for AstNode<'a, TSConditionalType<'a>> {
     }
 }
 
-impl<'a> FormatWrite<'a> for AstNode<'a, TSUnionType<'a>> {
-    fn write(&self, f: &mut Formatter<'_, 'a>) -> FormatResult<()> {
-        let mut types = self.types().iter();
-        if self.needs_parentheses(f) {
-            write!(f, "(")?;
-        }
-        if let Some(item) = types.next() {
-            write!(f, item)?;
-
-            for item in types {
-                write!(f, [" | ", item])?;
-            }
-        }
-        if self.needs_parentheses(f) {
-            write!(f, ")")?;
-        }
-        Ok(())
-    }
-}
-
 impl<'a> FormatWrite<'a> for AstNode<'a, TSIntersectionType<'a>> {
     fn write(&self, f: &mut Formatter<'_, 'a>) -> FormatResult<()> {
         let mut types = self.types().iter();

crates/oxc_formatter/src/write/union_type.rs

@@ -0,0 +1,224 @@
+use oxc_allocator::Vec;
+use oxc_ast::ast::*;
+use oxc_span::GetSpan;
+
+use crate::{
+    format_args,
+    formatter::{FormatResult, Formatter, prelude::*, trivia::FormatTrailingComments},
+    generated::ast_nodes::{AstNode, AstNodes},
+    parentheses::NeedsParentheses,
+    write,
+    write::FormatWrite,
+};
+
+impl<'a> FormatWrite<'a> for AstNode<'a, TSUnionType<'a>> {
+    fn write(&self, f: &mut Formatter<'_, 'a>) -> FormatResult<()> {
+        let types = self.types();
+
+        if types.len() == 1 {
+            return write!(f, self.types().first());
+        }
+
+        // ```ts
+        // {
+        //   a: string
+        // } | null | void
+        // ```
+        // should be inlined and not be printed in the multi-line variant
+        let should_hug = should_hug_type(self);
+        if should_hug {
+            return format_union_types(self.types(), true, f);
+        }
+
+        // Find the head of the nest union type chain
+        // ```js
+        // type Foo = | (| (A | B))
+        //                  ^^^^^
+        // ```
+        // If the current union type is `A | B`
+        // - `A | B` is the inner union type of `| (A | B)`
+        // - `| (A | B)` is the inner union type of `| (| (A | B))`
+        //
+        // So the head of the current nested union type chain is `| (| (A | B))`
+        // if we encounter a leading comment when navigating up the chain,
+        // we consider the current union type as having leading comments
+        let mut has_leading_comments = f.comments().has_comment_before(self.span().start);
+        let mut union_type_at_top = self;
+
+        while let AstNodes::TSUnionType(parent) = union_type_at_top.parent {
+            if parent.types().len() == 1 {
+                if f.comments().has_comment_before(parent.span().start) {
+                    has_leading_comments = true;
+                }
+                union_type_at_top = parent;
+            } else {
+                break;
+            }
+        }
+
+        let should_indent = {
+            let parent = union_type_at_top.parent;
+
+            // These parents have indent for their content, so we don't need to indent here
+            !match parent {
+                AstNodes::TSTypeAliasDeclaration(_) => has_leading_comments,
+                AstNodes::TSTypeAssertion(_)
+                | AstNodes::TSTupleType(_)
+                | AstNodes::TSTypeParameterInstantiation(_) => true,
+                _ => false,
+            }
+        };
+
+        let types = format_with(|f| {
+            if has_leading_comments {
+                write!(f, [soft_line_break()])?;
+            }
+
+            let leading_soft_line_break_or_space = should_indent && !has_leading_comments;
+
+            let separator = format_with(|f| {
+                if leading_soft_line_break_or_space {
+                    write!(f, [soft_line_break_or_space()])?;
+                }
+                write!(f, [text("|"), space()])
+            });
+
+            write!(f, [if_group_breaks(&separator)])?;
+
+            format_union_types(types, false, f)
+        });
+
+        let content = format_with(|f| {
+            // it is necessary to add parentheses for unions in intersections
+            // ```ts
+            // type Some = B & (C | A) & D
+            // ```
+            if self.needs_parentheses(f) {
+                return write!(f, [indent(&types), soft_line_break()]);
+            }
+
+            let is_inside_complex_tuple_type = match self.parent {
+                AstNodes::TSTupleType(tuple) => tuple.element_types().len() > 1,
+                _ => false,
+            };
+
+            if is_inside_complex_tuple_type {
+                write!(
+                    f,
+                    [
+                        indent(&format_args!(
+                            if_group_breaks(&format_args!(text("("), soft_line_break())),
+                            types
+                        )),
+                        soft_line_break(),
+                        if_group_breaks(&text(")"))
+                    ]
+                )
+            } else if should_indent {
+                write!(f, [indent(&types)])
+            } else {
+                write!(f, [types])
+            }
+        });
+
+        write!(f, [group(&content)])
+    }
+}
+
+fn should_hug_type(node: &AstNode<'_, TSUnionType<'_>>) -> bool {
+    // Simple heuristic: hug unions with object types and simple nullable types
+    let types = node.types();
+
+    if types.len() <= 3 {
+        let has_object_type = types.iter().any(|t| matches!(t.as_ref(), TSType::TSTypeLiteral(_)));
+
+        let has_simple_types = types.iter().any(|t| {
+            matches!(
+                t.as_ref(),
+                TSType::TSNullKeyword(_) | TSType::TSUndefinedKeyword(_) | TSType::TSVoidKeyword(_)
+            )
+        });
+
+        return has_object_type && has_simple_types;
+    }
+
+    false
+}
+
+pub struct FormatTSType<'a, 'b> {
+    next_node_span: Option<Span>,
+    element: &'b AstNode<'a, TSType<'a>>,
+    should_hug: bool,
+}
+
+impl<'a> Format<'a> for FormatTSType<'a, '_> {
+    fn fmt(&self, f: &mut crate::formatter::Formatter<'_, 'a>) -> FormatResult<()> {
+        let format_element = format_once(|f| {
+            self.element.fmt(f)?;
+            Ok(())
+        });
+        if self.should_hug {
+            write!(f, [format_element])?;
+        } else {
+            write!(f, [align(2, &format_element)])?;
+        }
+
+        if let Some(next_node_span) = self.next_node_span {
+            let comments_before_separator =
+                f.context().comments().comments_before_character(self.element.span().end, b'|');
+            FormatTrailingComments::Comments(comments_before_separator).fmt(f)?;
+
+            // ```ts
+            // type Some = A |
+            // // comment
+            // B
+            // ```
+            // to
+            // ```ts
+            // type Some =
+            // | A
+            // // comment
+            // | B
+            // ```
+            // If there is a leading own line comment between `|` and the next node, we need to put print comments
+            // before `|` instead of after it.
+            if f.comments().has_leading_own_line_comment(next_node_span.start) {
+                let comments = f.context().comments().comments_before(next_node_span.start);
+                FormatTrailingComments::Comments(comments).fmt(f)?;
+            }
+
+            if self.should_hug {
+                write!(f, [space()])?;
+            } else {
+                write!(f, [soft_line_break_or_space()])?;
+            }
+            write!(f, ["|"])
+        } else {
+            // ```ts
+            // type Foo = (
+            // | "thing1" // comment1
+            // | "thing2" // comment2
+            //            ^^^^^^^^^^^ the following logic is to print comment2,
+            // )[]; // comment 3
+            //```
+            // TODO: We may need to tweak `AstNode<'a, Vec<'a, T>>` iterator as some of Vec's last elements should have the following span.
+            let comments =
+                f.context().comments().end_of_line_comments_after(self.element.span().end);
+            FormatTrailingComments::Comments(comments).fmt(f)
+        }
+    }
+}
+
+fn format_union_types<'a>(
+    node: &AstNode<'a, Vec<'a, TSType<'a>>>,
+    should_hug: bool,
+    f: &mut Formatter<'_, 'a>,
+) -> FormatResult<()> {
+    f.join_with(space())
+        .entries(node.iter().enumerate().map(|(index, item)| FormatTSType {
+            next_node_span: node.get(index + 1).map(GetSpan::span),
+            element: item,
+            should_hug,
+        }))
+        .finish()
+}

tasks/ast_tools/src/generators/formatter/format.rs

@@ -38,7 +38,7 @@ const AST_NODE_WITHOUT_PRINTING_COMMENTS_LIST: &[&str] = &[
 ];
 
 const AST_NODE_NEEDS_PARENTHESES: &[&str] =
-    &["TSTypeAssertion", "TSInferType", "TSConditionalType"];
+    &["TSTypeAssertion", "TSInferType", "TSConditionalType", "TSUnionType"];
 
 const NEEDS_IMPLEMENTING_FMT_WITH_OPTIONS: phf::Map<&'static str, &'static str> = phf::phf_map! {
     "ArrowFunctionExpression" => "FormatJsArrowFunctionExpressionOptions",