Improve functionality and improve tests

package.json

@@ -17,18 +17,18 @@
     "index.js"
   ],
   "devDependencies": {
-    "@types/hast": "^2.3.4",
-    "hastscript": "^7.2.0",
-    "prettier": "^2.8.8",
-    "rehype-format": "^4.0.1",
-    "rehype-parse": "^8.0.4",
-    "rehype-stringify": "^9.0.3",
-    "typescript": "^5.0.4",
-    "unified": "^10.1.2"
+    "@types/hast": "^3.0.4",
+    "hastscript": "^9.0.0",
+    "prettier": "^3.2.5",
+    "rehype-format": "^5.0.0",
+    "rehype-parse": "^9.0.0",
+    "rehype-stringify": "^10.0.0",
+    "typescript": "^5.3.3",
+    "unified": "^11.0.4"
   },
   "dependencies": {
-    "hast-util-is-element": "^2.1.3",
-    "unist-util-visit": "^4.1.2"
+    "hast-util-is-element": "^3.0.0",
+    "unist-util-visit": "^5.0.0"
   },
   "prettier": {
     "proseWrap": "always",

src/components.ts

@@ -1,19 +1,23 @@
-import { Root, Properties, ElementContent } from "hast";
-import { Plugin, Processor } from "unified";
+import type { Root, Properties, ElementContent } from "hast";
+import type { Plugin, Processor } from "unified";
 import { visit, SKIP } from "unist-util-visit";
 import { isElement } from "hast-util-is-element";
-import { VFile } from "vfile";
+import type { VFile } from "vfile";
 
 export type ComponentFunction = (
   props: Properties,
   children: ElementContent[],
-  context: ComponentContext
-) => ElementContent;
+  context: ComponentContext,
+) => ElementContent | ElementContent[] | undefined | null;
 
+/**
+ * Context object passed to the component function to give it access to the
+ * current root of the tree, the current vfile and the processor.
+ */
 export interface ComponentContext {
   tree: Root;
   vfile: VFile;
-  processor: any;
+  processor: Processor;
 }
 
 interface Options {
@@ -22,29 +26,70 @@ interface Options {
 
 const rehypeComponents: Plugin<[Options], Root, Root> = function (options) {
   const { components = {} } = options;
-  const processor = this;
   return (tree, vfile) => {
-    const context: ComponentContext = { tree, vfile, processor };
+    const context: ComponentContext = { tree, vfile, processor: this };
     visit(tree, (node, index, parent) => {
       if (!isElement(node)) {
         return;
       }
-      node.properties;
       const component = components[node.tagName];
-      if (component && parent && index !== null) {
-        const replacedNode = component(
+      if (component && parent && index != null) {
+        // Invoke the component function
+        let returnedContent = component(
           node.properties || {},
           node.children,
-          context
+          context,
         );
-        parent.children[index] = replacedNode;
 
-        // This return value makes sure that the traversal continues by
-        // visiting the children of the replaced node (if any)
+        // Normalize returned content to an array
+        if (returnedContent == null) {
+          returnedContent = [];
+        } else if (!Array.isArray(returnedContent)) {
+          returnedContent = [returnedContent];
+        }
+
+        if (!returnedContent.every(isElementContent)) {
+          throw new Error(
+            `rehype-components: Component function is expected to return ElementContent or an array of ElementContent, but got ${JSON.stringify(
+              returnedContent,
+            )}.`,
+          );
+        }
+
+        // Replace the node with the normalized content
+        parent.children.splice(index, 1, ...returnedContent);
+
+        // This return value is a tuple that tells unist-util-visit to skip the
+        // children of the replaced node. Because we may have replaced the node.
         return [SKIP, index];
       }
     });
   };
 };
 
+function isElementContent(value: unknown): value is ElementContent {
+  if (isElement(value)) {
+    return true;
+  }
+
+  if (objectHasKey(value, "type")) {
+    if (value.type === "text") {
+      return true;
+    }
+
+    if (value.type === "comment") {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+function objectHasKey<T extends string>(
+  obj: unknown,
+  key: T,
+): obj is { [K in T]: unknown } {
+  return typeof obj === "object" && obj !== null && key in obj;
+}
+
 export default rehypeComponents;

test.js

@@ -1,37 +1,37 @@
+import { test } from "node:test";
+import assert from "node:assert/strict";
+
 import { unified } from "unified";
 import rehypeParse from "rehype-parse";
 import rehypeStringify from "rehype-stringify";
 import rehypeFormat from "rehype-format";
 import { h } from "hastscript";
-import { test } from "node:test";
-import assert from "node:assert/strict";
 
 import rehypeComponents from "./dist/components.js";
 
-function processDocument(document) {
-  return unified()
+function processDocument(document, options, useFormat = false) {
+  let processor = unified()
     .use(rehypeParse, { fragment: true })
-    .use(rehypeComponents, {
-      components: {
-        "documentation-page": DocumentationPage,
-        "info-box": InfoBox,
-        "copyright-notice": CopyrightNotice,
-      },
-    })
-    .use(rehypeStringify)
-    .use(rehypeFormat)
-    .process(document);
+    .use(rehypeComponents, options)
+    .use(rehypeStringify);
+
+  if (useFormat) {
+    processor = processor.use(rehypeFormat);
+  }
+
+  return processor.process(document);
 }
 
-const input = `
+test("example from readme", async () => {
+  const input = `
 <documentation-page title="Welcome">
   <info-box title="Reminder">Don't forget to run npm install</info-box>
   <p>Lorem ipsum...</p>
   <copyright-notice year="2020"></copyright-notice>
 </documentation-page>
 `;
 
-const expected = `
+  const expected = `
 <article class="documentation">
   <h1>Welcome</h1>
   <div class="infobox">
@@ -43,20 +43,104 @@ const expected = `
 </article>
 `;
 
-const DocumentationPage = (properties, children) =>
-  h("article.documentation", [h("h1", properties.title), ...children]);
+  const DocumentationPage = (properties, children) =>
+    h("article.documentation", [h("h1", properties.title), ...children]);
 
-const CopyrightNotice = (properties, children) =>
-  h("footer.notice", `© ${properties.year}`);
+  const CopyrightNotice = properties =>
+    h("footer.notice", `© ${properties.year}`);
 
-const InfoBox = (properties, children, context) =>
-  h(
-    ".infobox",
-    h(".infobox-title", properties.title || "Info"),
-    h(".infobox-body", children)
-  );
+  const InfoBox = (properties, children) =>
+    h(
+      ".infobox",
+      h(".infobox-title", properties.title || "Info"),
+      h(".infobox-body", children)
+    );
 
-test("example", async () => {
-  const output = await processDocument(input);
+  const output = await processDocument(
+    input,
+    {
+      components: {
+        "documentation-page": DocumentationPage,
+        "info-box": InfoBox,
+        "copyright-notice": CopyrightNotice,
+      },
+    },
+    true
+  );
   assert.deepEqual(String(output), expected);
 });
+
+test("component returns an array of nodes", async () => {
+  const input = `<wrapper><multiple-nodes></multiple-nodes></wrapper>`;
+  const expected = `<div><span>First</span><span>Second</span></div>`;
+
+  const wrapper = (properties, children) => h("div", children);
+  const multipleNodes = () => [h("span", "First"), h("span", "Second")];
+
+  const output = await processDocument(input, {
+    components: {
+      wrapper: wrapper,
+      "multiple-nodes": multipleNodes,
+    },
+  });
+  assert.strictEqual(String(output), expected);
+});
+
+test("component returns undefined or an empty array", async () => {
+  const input = `<wrapper><empty-component></empty-component><undefined-component></undefined-component></wrapper>`;
+  const expected = `<div></div>`; // Expect wrapper with no children
+
+  const wrapper = (properties, children) => h("div", children);
+  const emptyComponent = () => [];
+  const undefinedComponent = () => undefined;
+
+  const output = await processDocument(input, {
+    components: {
+      wrapper: wrapper,
+      "empty-component": emptyComponent,
+      "undefined-component": undefinedComponent,
+    },
+  });
+  assert.strictEqual(String(output), expected);
+});
+
+test("component returns invalid content", async () => {
+  const input = `<wrapper><invalid-content></invalid-content></wrapper>`;
+
+  const wrapper = (properties, children) => h("div", children);
+  const invalidContent = () => "This is invalid"; // Should cause an error
+
+  await assert.rejects(
+    processDocument(input, {
+      components: {
+        wrapper: wrapper,
+        "invalid-content": invalidContent,
+      },
+    }),
+    err => {
+      assert.strictEqual(
+        err.message.includes("expected to return ElementContent"),
+        true
+      );
+      return true;
+    },
+    "Did not throw with expected error message"
+  );
+});
+
+test("component returns a tree with another component for recursive rendering", async () => {
+  const input = `<parent-component></parent-component>`;
+  const expected = `<div class="parent"><article class="child">Child content</article></div>`;
+
+  const ParentComponent = () => h("div.parent", h("child-component"));
+
+  const ChildComponent = () => h("article.child", "Child content");
+
+  const output = await processDocument(input, {
+    components: {
+      "parent-component": ParentComponent,
+      "child-component": ChildComponent,
+    },
+  });
+  assert.strictEqual(String(output), expected);
+});

tsconfig.json

@@ -3,8 +3,9 @@
     "target": "es2020",
     "strict": true,
     "outDir": "./dist",
-    "module": "ES2020",
-    "moduleResolution": "nodenext"
+    "module": "nodenext",
+    "moduleResolution": "nodenext",
+    "declaration": true
   },
   "exclude": ["node_modules", "dist"]
 }