feat(mdx-loader): the table-of-contents should display toc/headings of imported MDX partials

packages/docusaurus-mdx-loader/src/remark/toc/__tests__/__snapshots__/index.test.ts.snap

@@ -236,3 +236,23 @@ Some content here
 export const c = 1;
 "
 `;
+
+exports[`toc remark plugin works with imported markdown 1`] = `
+"import Partial from './_partial.md';
+
+export const toc = [
+	{
+		value: 'Thanos',
+		id: 'thanos',
+		level: 2
+	},
+	...toc0
+]
+
+## Thanos
+
+Foo
+
+<Partial />
+"
+`;

packages/docusaurus-mdx-loader/src/remark/toc/__tests__/index.test.ts

@@ -70,4 +70,9 @@ describe('toc remark plugin', () => {
     const result = await processFixture('empty-headings');
     expect(result).toMatchSnapshot();
   });
+
+  it('works with imported markdown', async () => {
+    const result = await processFixture('imported-markdown');
+    expect(result).toMatchSnapshot();
+  });
 });

packages/docusaurus-mdx-loader/src/remark/toc/index.ts

@@ -7,15 +7,15 @@
 
 import {parse, type ParserOptions} from '@babel/parser';
 import traverse from '@babel/traverse';
-import stringifyObject from 'stringify-object';
-import {toValue} from '../utils';
+import {constructArrayString, toValue} from '../utils';
 import type {Identifier} from '@babel/types';
 import type {Node, Parent} from 'unist';
 import type {Heading, Literal} from 'mdast';
 // @ts-expect-error: TODO see https://github.com/microsoft/TypeScript/issues/49721
 import type {Transformer} from 'unified';
 import type {
   MdxjsEsm,
+  MdxJsxFlowElement,
   // @ts-expect-error: TODO see https://github.com/microsoft/TypeScript/issues/49721
 } from 'mdast-util-mdx';
 
@@ -93,21 +93,75 @@ const plugin: Plugin = function plugin(
     const {toString} = await import('mdast-util-to-string');
     const {visit} = await import('unist-util-visit');
 
-    const headings: TOCItem[] = [];
+    const partialComponentToHeadingsName: {[key: string]: string} =
+      Object.create(null);
 
-    visit(root, 'heading', (child: Heading) => {
-      const value = toString(child);
+    // TOCItem or string already with the spread operator
+    const headings: (TOCItem | string)[] = [];
 
-      // depth:1 headings are titles and not included in the TOC
-      if (!value || child.depth < 2) {
-        return;
+    visit(root, ['heading', 'mdxjsEsm', 'mdxJsxFlowElement'], (child) => {
+      if (child.type === 'heading') {
+        const headingNode = child as Heading;
+        const value = toString(headingNode);
+
+        // depth:1 headings are titles and not included in the TOC
+        if (!value || headingNode.depth < 2) {
+          return;
+        }
+
+        headings.push({
+          value: toValue(headingNode, toString),
+          id: headingNode.data!.id!,
+          level: headingNode.depth,
+        });
       }
 
-      headings.push({
-        value: toValue(child, toString),
-        id: child.data!.id!,
-        level: child.depth,
-      });
+      if (child.type === 'mdxjsEsm') {
+        const importNode = child as MdxjsEsm;
+        if (!importNode.data?.estree) {
+          return;
+        }
+
+        for (const importDeclaration of importNode.data.estree.body) {
+          if (importDeclaration.type !== 'ImportDeclaration') {
+            continue;
+          }
+          const importPath = importDeclaration.source.value as string;
+          const isMdxImport = /\.mdx?$/.test(importPath);
+          if (!isMdxImport) {
+            continue;
+          }
+
+          const componentName = importDeclaration.specifiers.find(
+            (o: any) => o.type === 'ImportDefaultSpecifier',
+          )?.local.name;
+
+          if (!componentName) {
+            continue;
+          }
+          const {length} = Object.keys(partialComponentToHeadingsName);
+          const exportAsName = `${name}${length}`;
+          partialComponentToHeadingsName[componentName] = exportAsName;
+
+          importDeclaration.specifiers.push({
+            type: 'ImportSpecifier',
+            imported: {type: 'Identifier', name},
+            local: {type: 'Identifier', name: exportAsName},
+          });
+        }
+      }
+
+      if (child.type === 'mdxJsxFlowElement') {
+        const node = child as MdxJsxFlowElement;
+        const nodeName = node.name;
+        if (!nodeName) {
+          return;
+        }
+        const headingsName = partialComponentToHeadingsName[nodeName];
+        if (headingsName) {
+          headings.push(`...${headingsName}`);
+        }
+      }
     });
 
     const {children} = root as Parent;
@@ -124,9 +178,21 @@ export default plugin;
 async function createExportNode(name: string, object: any): Promise<MdxjsEsm> {
   const {valueToEstree} = await import('estree-util-value-to-estree');
 
+  const tocObject = object.map((heading: TOCItem | string) => {
+    if (typeof heading === 'string') {
+      const argumentName = heading.replace('...', '');
+      return {
+        type: 'SpreadElement',
+        argument: {type: 'Identifier', name: argumentName},
+      };
+    }
+
+    return valueToEstree(heading);
+  });
+
   return {
     type: 'mdxjsEsm',
-    value: `export const ${name} = ${stringifyObject(object)}`,
+    value: `export const ${name} = ${constructArrayString(object)}`,
     data: {
       estree: {
         type: 'Program',
@@ -142,7 +208,10 @@ async function createExportNode(name: string, object: any): Promise<MdxjsEsm> {
                     type: 'Identifier',
                     name,
                   },
-                  init: valueToEstree(object),
+                  init: {
+                    type: 'ArrayExpression',
+                    elements: tocObject,
+                  },
                 },
               ],
               kind: 'const',

packages/docusaurus-mdx-loader/src/remark/utils/index.ts

@@ -6,6 +6,7 @@
  */
 
 import escapeHtml from 'escape-html';
+import stringifyObject from 'stringify-object';
 import type {Parent, Node} from 'unist';
 import type {PhrasingContent, Heading} from 'mdast';
 import type {
@@ -154,3 +155,32 @@ export function assetRequireAttributeValue(
     },
   };
 }
+
+/**
+ * Similar to stringify-object, but keeps spread operators,
+ * instead of turning them into strings.
+ * @param objects
+ */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function constructArrayString(objects: any[]): string {
+  let result = '[';
+
+  for (const obj of objects) {
+    if (typeof obj === 'string') {
+      result = `${result}\n\t${obj},`;
+    } else {
+      result = `${result}\n\t${stringifyObject(obj).replace(/\n/g, '\n\t')},`;
+    }
+  }
+
+  // Remove trailing coma
+  result = result.replace(/,$/, '');
+  // Close array on same line if it is empty
+  if (result === '[') {
+    result += ']';
+  } else {
+    result += '\n]';
+  }
+
+  return result;
+}

website/_dogfooding/_pages tests/_anotherPagePartial.mdx

@@ -0,0 +1,7 @@
+### Another page partial content
+
+This is text coming from another page partial
+
+#### Foo
+
+Level 4 headings don't belong in ToC

website/_dogfooding/_pages tests/index.mdx

@@ -36,6 +36,7 @@ import Readme from "../README.mdx"
 - [Tabs tests](/tests/pages/tabs-tests)
 - [z-index tests](/tests/pages/z-index-tests)
 - [Head metadata tests](/tests/pages/head-metadata)
+- [Partials tests](/tests/pages/partials-tests)
 - [Unlisted page](/tests/pages/unlisted)
 - [Analytics](/tests/pages/analytics)
 - [Embeds](/tests/pages/embeds)

website/_dogfooding/_pages tests/partials-tests.mdx

@@ -0,0 +1,12 @@
+import PagePartial from './_pagePartial.mdx';
+import AnotherPagePartial from './_anotherPagePartial.mdx';
+
+# Partials tests
+
+This page consists of multiple files imported into one. Notice how the table of contents works even for imported headings.
+
+## Imported content
+
+<PagePartial />
+
+<AnotherPagePartial />