Fix Sandpack MDX parsing and typegen

Author: feedthejim

next-env.d.ts

@@ -1,7 +1,7 @@
 /// <reference types="next" />
 /// <reference types="next/image-types/global" />
 /// <reference types="next/navigation-types/compat/navigation" />
-import './.next/dev/types/routes.d.ts';
+import './.next/types/routes.d.ts';
 
 // NOTE: This file should not be edited
 // see https://nextjs.org/docs/app/api-reference/config/typescript for more information.

src/components/MDX/ExpandableExample.tsx

@@ -28,6 +28,11 @@ interface ExpandableExampleProps {
   type: 'DeepDive' | 'Example';
 }
 
+type ExpandableTitleElement = React.ReactElement<{
+  id?: string;
+  children?: React.ReactNode;
+}>;
+
 function getExpandableChildren(children: React.ReactNode) {
   return Children.toArray(children).filter((child) => {
     return !(typeof child === 'string' && child.trim() === '');
@@ -50,10 +55,11 @@ function ExpandableExample({children, excerpt, type}: ExpandableExampleProps) {
       `Expandable content ${type} is missing a corresponding title at the beginning`
     );
   }
+  const titleElement = titleChild as ExpandableTitleElement;
 
   const isDeepDive = type === 'DeepDive';
   const isExample = type === 'Example';
-  const id = titleChild.props.id;
+  const id = titleElement.props.id;
   const hash = useLocationHash();
   const [isExpanded, setIsExpanded] = useState(false);
   const [isAutoExpandedDismissed, setIsAutoExpandedDismissed] = useState(false);
@@ -105,7 +111,7 @@ function ExpandableExample({children, excerpt, type}: ExpandableExampleProps) {
           <H4
             id={id}
             className="text-xl font-bold text-primary dark:text-primary-dark">
-            {titleChild.props.children}
+            {titleElement.props.children}
           </H4>
           {excerpt && <div>{excerpt}</div>}
         </div>

src/components/MDX/Sandpack/SandpackRSCRoot.tsx

@@ -1,3 +1,5 @@
+'use client';
+
 /**
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
@@ -9,18 +11,18 @@
  * Copyright (c) Facebook, Inc. and its affiliates.
  */
 
-import {Children} from 'react';
 import * as React from 'react';
 import {SandpackProvider} from '@codesandbox/sandpack-react/unstyled';
 import {SandpackLogLevel} from '@codesandbox/sandpack-client';
 import {CustomPreset} from './CustomPreset';
-import {createFileMap} from './createFileMap';
+import type {SandpackSnippetFile} from './createFileMap';
 import {CustomTheme} from './Themes';
 import {templateRSC} from './templateRSC';
 import {RscFileBridge} from './sandpack-rsc/RscFileBridge';
 
 type SandpackProps = {
-  children: React.ReactNode;
+  files: Record<string, SandpackSnippetFile>;
+  providedFiles: string[];
   autorun?: boolean;
 };
 
@@ -75,9 +77,7 @@ ul {
 `.trim();
 
 function SandpackRSCRoot(props: SandpackProps) {
-  const {children, autorun = true} = props;
-  const codeSnippets = Children.toArray(children) as React.ReactElement[];
-  const files = createFileMap(codeSnippets);
+  const {files, providedFiles, autorun = true} = props;
 
   if ('/index.html' in files) {
     throw new Error(
@@ -86,15 +86,18 @@ function SandpackRSCRoot(props: SandpackProps) {
     );
   }
 
-  files['/src/styles.css'] = {
-    code: [sandboxStyle, files['/src/styles.css']?.code ?? ''].join('\n\n'),
-    hidden: !files['/src/styles.css']?.visible,
+  const sandpackFiles = {
+    ...files,
+    '/src/styles.css': {
+      code: [sandboxStyle, files['/src/styles.css']?.code ?? ''].join('\n\n'),
+      hidden: !files['/src/styles.css']?.visible,
+    },
   };
 
   return (
     <div className="sandpack sandpack--playground w-full my-8" dir="ltr">
       <SandpackProvider
-        files={{...templateRSC, ...files}}
+        files={{...templateRSC, ...sandpackFiles}}
         theme={CustomTheme}
         customSetup={{
           dependencies: {},
@@ -108,7 +111,7 @@ function SandpackRSCRoot(props: SandpackProps) {
         }}>
         <RscFileBridge />
         <CustomPreset
-          providedFiles={Object.keys(files)}
+          providedFiles={providedFiles}
           showOpenInCodeSandbox={false}
         />
       </SandpackProvider>

src/components/MDX/Sandpack/SandpackRoot.tsx

@@ -1,3 +1,5 @@
+'use client';
+
 /**
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
@@ -9,17 +11,17 @@
  * Copyright (c) Facebook, Inc. and its affiliates.
  */
 
-import {Children} from 'react';
 import * as React from 'react';
 import {SandpackProvider} from '@codesandbox/sandpack-react/unstyled';
 import {SandpackLogLevel} from '@codesandbox/sandpack-client';
 import {CustomPreset} from './CustomPreset';
-import {createFileMap} from './createFileMap';
+import type {SandpackSnippetFile} from './createFileMap';
 import {CustomTheme} from './Themes';
 import {template} from './template';
 
 type SandpackProps = {
-  children: React.ReactNode;
+  files: Record<string, SandpackSnippetFile>;
+  providedFiles: string[];
   autorun?: boolean;
 };
 
@@ -74,9 +76,7 @@ ul {
 `.trim();
 
 function SandpackRoot(props: SandpackProps) {
-  let {children, autorun = true} = props;
-  const codeSnippets = Children.toArray(children) as React.ReactElement[];
-  const files = createFileMap(codeSnippets);
+  const {files, providedFiles, autorun = true} = props;
 
   if ('/index.html' in files) {
     throw new Error(
@@ -85,15 +85,18 @@ function SandpackRoot(props: SandpackProps) {
     );
   }
 
-  files['/src/styles.css'] = {
-    code: [sandboxStyle, files['/src/styles.css']?.code ?? ''].join('\n\n'),
-    hidden: !files['/src/styles.css']?.visible,
+  const sandpackFiles = {
+    ...files,
+    '/src/styles.css': {
+      code: [sandboxStyle, files['/src/styles.css']?.code ?? ''].join('\n\n'),
+      hidden: !files['/src/styles.css']?.visible,
+    },
   };
 
   return (
     <div className="sandpack sandpack--playground w-full my-8" dir="ltr">
       <SandpackProvider
-        files={{...template, ...files}}
+        files={{...template, ...sandpackFiles}}
         theme={CustomTheme}
         customSetup={{
           environment: 'create-react-app',
@@ -105,7 +108,7 @@ function SandpackRoot(props: SandpackProps) {
           bundlerURL: 'https://786946de.sandpack-bundler-4bw.pages.dev',
           logLevel: SandpackLogLevel.None,
         }}>
-        <CustomPreset providedFiles={Object.keys(files)} />
+        <CustomPreset providedFiles={providedFiles} />
       </SandpackProvider>
     </div>
   );

src/components/MDX/Sandpack/createFileMap.ts

@@ -10,13 +10,23 @@
  */
 
 import type {SandpackFile} from '@codesandbox/sandpack-react/unstyled';
-import type {PropsWithChildren, ReactElement, HTMLAttributes} from 'react';
+import {Children, isValidElement} from 'react';
+import type {
+  PropsWithChildren,
+  ReactElement,
+  HTMLAttributes,
+  ReactNode,
+} from 'react';
 import {getMDXName} from '../getMDXName';
 
 export const AppJSPath = `/src/App.js`;
 export const StylesCSSPath = `/src/styles.css`;
 export const SUPPORTED_FILES = [AppJSPath, StylesCSSPath];
 
+export type SandpackSnippetFile = SandpackFile & {
+  visible?: boolean;
+};
+
 /**
  * Tokenize meta attributes while ignoring brace-wrapped metadata (e.g. {expectedErrors: …}).
  */
@@ -77,22 +87,105 @@ function splitMeta(meta: string): string[] {
   return tokens;
 }
 
+function collectCodeSnippets(children: ReactNode): ReactElement[] {
+  const codeSnippets: ReactElement[] = [];
+
+  Children.forEach(children, (child) => {
+    if (!isValidElement(child)) {
+      return;
+    }
+
+    if (getMDXName(child) === 'pre') {
+      codeSnippets.push(child);
+      return;
+    }
+
+    const props = child.props as {children?: ReactNode} | null;
+    if (props?.children != null) {
+      codeSnippets.push(...collectCodeSnippets(props.children));
+    }
+  });
+
+  return codeSnippets;
+}
+
+type CodeElementProps = HTMLAttributes<HTMLDivElement> & {
+  meta?: string;
+  children?: ReactNode;
+};
+
+function findCodeElement(
+  children: ReactNode
+): ReactElement<CodeElementProps> | null {
+  let codeElement: ReactElement<CodeElementProps> | null = null;
+
+  Children.forEach(children, (child) => {
+    if (codeElement || !isValidElement(child)) {
+      return;
+    }
+
+    const props = child.props as CodeElementProps | null;
+    if (
+      getMDXName(child) === 'code' ||
+      typeof props?.meta === 'string' ||
+      (typeof props?.className === 'string' &&
+        props.className.startsWith('language-'))
+    ) {
+      codeElement = child as ReactElement<CodeElementProps>;
+      return;
+    }
+
+    if (props?.children != null) {
+      codeElement = findCodeElement(props.children);
+    }
+  });
+
+  return codeElement;
+}
+
+function getTextContent(node: ReactNode): string {
+  let text = '';
+
+  Children.forEach(node, (child) => {
+    if (typeof child === 'string' || typeof child === 'number') {
+      text += child;
+      return;
+    }
+
+    if (!isValidElement(child)) {
+      return;
+    }
+
+    const props = child.props as {children?: ReactNode} | null;
+    text += getTextContent(props?.children ?? null);
+  });
+
+  return text;
+}
+
 export const createFileMap = (codeSnippets: any) => {
-  return codeSnippets.reduce(
-    (result: Record<string, SandpackFile>, codeSnippet: React.ReactElement) => {
-      if (getMDXName(codeSnippet) !== 'pre') {
-        return result;
+  return collectCodeSnippets(codeSnippets).reduce(
+    (
+      result: Record<string, SandpackSnippetFile>,
+      codeSnippet: React.ReactElement
+    ) => {
+      const codeElement = findCodeElement(
+        (
+          codeSnippet.props as PropsWithChildren<{
+            children?: ReactNode;
+          }>
+        ).children
+      );
+
+      if (!codeElement) {
+        throw new Error('Code block is missing a code element.');
       }
-      const {props} = (
-        codeSnippet.props as PropsWithChildren<{
-          children: ReactElement<
-            HTMLAttributes<HTMLDivElement> & {meta?: string}
-          >;
-        }>
-      ).children;
+
+      const props = codeElement.props;
       let filePath; // path in the folder structure
       let fileHidden = false; // if the file is available as a tab
       let fileActive = false; // if the file tab is shown by default
+      let fileVisible = false; // if the file tab should be forced visible
 
       if (props.meta) {
         const tokens = splitMeta(props.meta);
@@ -108,6 +201,9 @@ export const createFileMap = (codeSnippets: any) => {
         if (tokens.includes('active')) {
           fileActive = true;
         }
+        if (tokens.includes('visible')) {
+          fileVisible = true;
+        }
       } else {
         if (props.className === 'language-js') {
           filePath = AppJSPath;
@@ -138,9 +234,10 @@ export const createFileMap = (codeSnippets: any) => {
         );
       }
       result[filePath] = {
-        code: (props.children || '') as string,
+        code: getTextContent(props.children ?? null),
         hidden: fileHidden,
         active: fileActive,
+        visible: fileVisible,
       };
 
       return result;