From da20f922e3525201c8681a9f6ea9ad86e042ad30 Mon Sep 17 00:00:00 2001
From: Vlad Moroz <hi@vladmoroz.com>
Date: Thu, 14 Nov 2024 17:52:10 +0100
Subject: [PATCH] Format union types

---
 docs/src/components/TableCode.tsx | 87 +++++++++++++++++++++++++++++++
 docs/src/mdx-components.tsx       |  3 +-
 2 files changed, 89 insertions(+), 1 deletion(-)
 create mode 100644 docs/src/components/TableCode.tsx

diff --git a/docs/src/components/TableCode.tsx b/docs/src/components/TableCode.tsx
new file mode 100644
index 000000000..f32e5eb21
--- /dev/null
+++ b/docs/src/components/TableCode.tsx
@@ -0,0 +1,87 @@
+import * as React from 'react';
+
+interface TableCodeProps extends React.ComponentProps<'code'> {
+  printWidth?: number;
+}
+
+/** An inline code component that breaks long union types into multiple lines */
+export function TableCode({ children, printWidth = 40, ...props }: TableCodeProps) {
+  const text = getTextContents(children);
+
+  if (text.includes('|') && text.length > printWidth) {
+    const unionGroups: React.ReactNode[][] = [];
+    const parts = React.Children.toArray(children);
+
+    let depth = 0;
+    let groupIndex = 0;
+    parts.forEach((child, index) => {
+      if (index === 0) {
+        unionGroups.push([child]);
+        return;
+      }
+
+      if (getTextContents(child).trim() === '|' && depth < 1) {
+        groupIndex += 1;
+        unionGroups.push([]);
+        return;
+      }
+
+      if (getTextContents(child).trim() === '(') {
+        depth += 1;
+      }
+
+      if (getTextContents(child).trim() === ')') {
+        depth -= 1;
+      }
+
+      unionGroups[groupIndex].push(child);
+    });
+
+    if (unionGroups.length > 1) {
+      unionGroups.forEach((_, index) => {
+        const pipe = <span style={{ color: 'var(--syntax-keyword)' }}>| </span>;
+        const pipeWithNewline = (
+          <React.Fragment>
+            <br />
+            {pipe}
+          </React.Fragment>
+        );
+
+        const element = index === 0 ? pipe : pipeWithNewline;
+        unionGroups.splice(index * 2, 0, [element]);
+      });
+    }
+
+    children = unionGroups.flat();
+  }
+
+  return <code {...props}>{children}</code>;
+}
+
+function getTextContents(node?: React.ReactNode): string {
+  if (hasChildren(node)) {
+    return getTextContents(node.props?.children);
+  }
+
+  if (Array.isArray(node)) {
+    return node.map(getTextContents).flat().filter(Boolean).join('');
+  }
+
+  if (typeof node === 'string') {
+    return node;
+  }
+
+  return '';
+}
+
+function hasChildren(
+  element?: React.ReactNode,
+): element is React.ReactElement<React.PropsWithChildren> {
+  return (
+    React.isValidElement(element) &&
+    typeof element.props === 'object' &&
+    !!element.props &&
+    'children' in element.props &&
+    Boolean(element.props.children)
+  );
+}
diff --git a/docs/src/mdx-components.tsx b/docs/src/mdx-components.tsx
index 59eea9d94..1c387b92a 100644
--- a/docs/src/mdx-components.tsx
+++ b/docs/src/mdx-components.tsx
@@ -7,6 +7,7 @@ import { Code } from './components/Code';
 import { PropsTable } from './components/reference/PropsTable';
 import { AttributesTable } from './components/reference/AttributesTable';
 import { CssVariablesTable } from './components/reference/CssVariablesTable';
+import { TableCode } from './components/TableCode';
 
 interface MDXComponents {
   [key: string]: React.FC<any> | MDXComponents;
@@ -69,7 +70,7 @@ export const tableMdxComponents: MDXComponents = {
   ...mdxComponents,
   // eslint-disable-next-line react/jsx-no-useless-fragment
   p: (props) => <React.Fragment {...props} />,
-  code: (props) => <code {...props} />,
+  code: TableCode,
 };
 
 export function useMDXComponents(): MDXComponents {