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

packages/docusaurus-mdx-loader/package.json

@@ -25,6 +25,7 @@
     "@docusaurus/utils-validation": "3.0.0",
     "@mdx-js/mdx": "^3.0.0",
     "@slorber/remark-comment": "^1.0.0",
+    "astring": "^1.8.6",
     "escape-html": "^1.0.3",
     "estree-util-value-to-estree": "^3.0.1",
     "file-loader": "^6.2.0",

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

@@ -236,3 +236,22 @@ Some content here
 export const c = 1;
 "
 `;
+
+exports[`toc remark plugin works with imported markdown 1`] = `
+"import Partial from './_partial.md';
+
+import {toc as toc0} 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,41 +7,41 @@
 
 import {parse, type ParserOptions} from '@babel/parser';
 import traverse from '@babel/traverse';
-import stringifyObject from 'stringify-object';
+import {generate} from 'astring';
 import {toValue} from '../utils';
+import {hasImports, isExport, isImport} from './utils';
+import type {TOCItem, NestedTOC} from './utils';
+import type {
+  SpreadElement,
+  Program,
+  ImportDeclaration,
+  ImportSpecifier,
+} from 'estree';
 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';
 
+// Reexport TOCItem, since it's used throughout the project
+export type {TOCItem} from './utils';
+
 // TODO as of April 2023, no way to import/re-export this ESM type easily :/
 // TODO upgrade to TS 5.3
 // See https://github.com/microsoft/TypeScript/issues/49721#issuecomment-1517839391
 // import type {Plugin} from 'unified';
-type Plugin = any; // TODO fix this asap
-
-export type TOCItem = {
-  readonly value: string;
-  readonly id: string;
-  readonly level: number;
-};
+type Plugin = any;
 
 const parseOptions: ParserOptions = {
   plugins: ['jsx'],
   sourceType: 'module',
 };
 
-const isImport = (child: any): child is Literal =>
-  child.type === 'mdxjsEsm' && child.value.startsWith('import');
-const hasImports = (index: number) => index > -1;
-const isExport = (child: any): child is Literal =>
-  child.type === 'mdxjsEsm' && child.value.startsWith('export');
-
 interface PluginOptions {
   name?: string;
 }
@@ -75,7 +75,7 @@ const getOrCreateExistingTargetIndex = async (
   });
 
   if (targetIndex === -1) {
-    const target = await createExportNode(name, []);
+    const target = await createExportNode(name, [], []);
 
     targetIndex = hasImports(importsIndex) ? importsIndex + 1 : 0;
     children.splice(targetIndex, 0, target);
@@ -93,66 +93,158 @@ 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);
+
+    const headings: (TOCItem | NestedTOC)[] = [];
+    const imports: ImportDeclaration[] = [];
 
-    visit(root, 'heading', (child: Heading) => {
-      const value = toString(child);
+    function visitHeading(node: Heading) {
+      const value = toString(node);
 
       // depth:1 headings are titles and not included in the TOC
-      if (!value || child.depth < 2) {
+      if (!value || node.depth < 2) {
         return;
       }
 
       headings.push({
-        value: toValue(child, toString),
-        id: child.data!.id!,
-        level: child.depth,
+        value: toValue(node, toString),
+        id: node.data!.id!,
+        level: node.depth,
       });
+    }
+
+    function visitMdxjsEsm(node: MdxjsEsm) {
+      if (!node.data?.estree) {
+        return;
+      }
+
+      for (const potentialImportDeclaration of node.data.estree.body) {
+        if (potentialImportDeclaration.type !== 'ImportDeclaration') {
+          continue;
+        }
+
+        const importPath = potentialImportDeclaration.source.value as string;
+        const isMdxImport = /\.mdx?$/.test(importPath);
+        if (!isMdxImport) {
+          continue;
+        }
+
+        const componentName = potentialImportDeclaration.specifiers.find(
+          (o: Node) => o.type === 'ImportDefaultSpecifier',
+        )?.local.name;
+
+        if (!componentName) {
+          continue;
+        }
+        const {length} = Object.keys(partialComponentToHeadingsName);
+        const exportAsName = `${name}${length}`;
+        partialComponentToHeadingsName[componentName] = exportAsName;
+
+        const specifier: ImportSpecifier = {
+          type: 'ImportSpecifier',
+          imported: {type: 'Identifier', name},
+          local: {type: 'Identifier', name: exportAsName},
+        };
+
+        imports.push({
+          type: 'ImportDeclaration',
+          specifiers: [specifier],
+          source: potentialImportDeclaration.source,
+        });
+        potentialImportDeclaration.specifiers.push(specifier);
+      }
+    }
+
+    function visitMdxJsxFlowElement(node: MdxJsxFlowElement) {
+      const nodeName = node.name;
+      if (!nodeName) {
+        return;
+      }
+      const headingsName = partialComponentToHeadingsName[nodeName];
+      if (headingsName) {
+        headings.push({
+          nested: true,
+          name: headingsName,
+        });
+      }
+    }
+
+    visit(root, ['heading', 'mdxjsEsm', 'mdxJsxFlowElement'], (child) => {
+      if (child.type === 'heading') {
+        visitHeading(child as Heading);
+      } else if (child.type === 'mdxjsEsm') {
+        visitMdxjsEsm(child as MdxjsEsm);
+      } else if (child.type === 'mdxJsxFlowElement') {
+        visitMdxJsxFlowElement(child as MdxJsxFlowElement);
+      }
     });
 
     const {children} = root as Parent;
     const targetIndex = await getOrCreateExistingTargetIndex(children, name);
 
     if (headings?.length) {
-      children[targetIndex] = await createExportNode(name, headings);
+      children[targetIndex] = await createExportNode(name, headings, imports);
     }
   };
 };
 
 export default plugin;
 
-async function createExportNode(name: string, object: any): Promise<MdxjsEsm> {
+async function createExportNode(
+  name: string,
+  headings: (TOCItem | NestedTOC)[],
+  imports: ImportDeclaration[],
+): Promise<MdxjsEsm> {
   const {valueToEstree} = await import('estree-util-value-to-estree');
 
+  const tocObject = headings.map((heading) => {
+    if ('nested' in heading) {
+      const spreadElement: SpreadElement = {
+        type: 'SpreadElement',
+        argument: {type: 'Identifier', name: heading.name},
+      };
+      return spreadElement;
+    }
+
+    return valueToEstree(heading);
+  });
+
+  const estree: Program = {
+    type: 'Program',
+    body: [
+      ...imports,
+      {
+        type: 'ExportNamedDeclaration',
+        declaration: {
+          type: 'VariableDeclaration',
+          declarations: [
+            {
+              type: 'VariableDeclarator',
+              id: {
+                type: 'Identifier',
+                name,
+              },
+              init: {
+                type: 'ArrayExpression',
+                elements: tocObject,
+              },
+            },
+          ],
+          kind: 'const',
+        },
+        specifiers: [],
+        source: null,
+      },
+    ],
+    sourceType: 'module',
+  };
+
   return {
     type: 'mdxjsEsm',
-    value: `export const ${name} = ${stringifyObject(object)}`,
+    value: generate(estree),
     data: {
-      estree: {
-        type: 'Program',
-        body: [
-          {
-            type: 'ExportNamedDeclaration',
-            declaration: {
-              type: 'VariableDeclaration',
-              declarations: [
-                {
-                  type: 'VariableDeclarator',
-                  id: {
-                    type: 'Identifier',
-                    name,
-                  },
-                  init: valueToEstree(object),
-                },
-              ],
-              kind: 'const',
-            },
-            specifiers: [],
-            source: null,
-          },
-        ],
-        sourceType: 'module',
-      },
+      estree,
     },
   };
 }

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

@@ -0,0 +1,40 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import type {Literal} from 'mdast';
+import type {Node} from 'unist';
+import type {
+  MdxjsEsm,
+  // @ts-expect-error: TODO see https://github.com/microsoft/TypeScript/issues/49721
+} from 'mdast-util-mdx';
+
+export type TOCItem = {
+  readonly value: string;
+  readonly id: string;
+  readonly level: number;
+};
+
+export type NestedTOC = {
+  readonly nested: true;
+  readonly name: string;
+};
+
+export const isImport = (child: Node): child is Literal => {
+  if (child.type === 'mdxjsEsm') {
+    return (child as MdxjsEsm).value.startsWith('import');
+  }
+  return false;
+};
+
+export const hasImports = (index: number): boolean => index > -1;
+
+export const isExport = (child: Node): child is Literal => {
+  if (child.type === 'mdxjsEsm') {
+    return (child as MdxjsEsm).value.startsWith('export');
+  }
+  return false;
+};

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 />

yarn.lock

@@ -4423,7 +4423,7 @@ astral-regex@^2.0.0:
   resolved "https://registry.yarnpkg.com/astral-regex/-/astral-regex-2.0.0.tgz#483143c567aeed4785759c0865786dc77d7d2e31"
   integrity sha512-Z7tMw1ytTXt5jqMcOP+OQteU1VuNK9Y02uuJtKQ1Sv69jXQKKg5cibLwGJow8yzZP+eAc18EmLGPal0bp36rvQ==
 
-astring@^1.8.0:
+astring@^1.8.0, astring@^1.8.6:
   version "1.8.6"
   resolved "https://registry.yarnpkg.com/astring/-/astring-1.8.6.tgz#2c9c157cf1739d67561c56ba896e6948f6b93731"
   integrity sha512-ISvCdHdlTDlH5IpxQJIex7BWBywFWgjJSVdwst+/iQCoEYnyOaQ95+X1JGshuBjGp6nxKUy1jMgE3zPqN7fQdg==