From 4d5e134893c0ec92ff7e4208edaf89fe52f28340 Mon Sep 17 00:00:00 2001
From: Titus Wormer <tituswormer@gmail.com>
Date: Wed, 18 Jan 2023 13:49:08 +0100
Subject: [PATCH] Add support for passing nodes to components

---
 lib/components.d.ts | 10 +++++++---
 lib/index.js        | 16 +++++++++++++---
 readme.md           | 28 ++++++++++++++++++----------
 test/index.js       | 33 +++++++++++++++++++++++++++++++++
 4 files changed, 71 insertions(+), 16 deletions(-)

diff --git a/lib/components.d.ts b/lib/components.d.ts
index e7abb16..2c8f2e1 100644
--- a/lib/components.d.ts
+++ b/lib/components.d.ts
@@ -1,3 +1,5 @@
+import type {Element} from 'hast'
+
 /**
  * Basic functional component: given props, returns an element.
  *
@@ -39,12 +41,14 @@ export type Component<ComponentProps> =
   | FunctionComponent<ComponentProps>
   | ClassComponent<ComponentProps>
 
+export type ExtraProps = {node?: Element | undefined}
+
 /**
  * Possible components to use.
  *
  * Each key is a tag name typed in `JSX.IntrinsicElements`.
- * Each value is a component accepting the corresponding props or a different
- * tag name.
+ * Each value is either a different tag name, or a component accepting the
+ * corresponding props (and an optional `node` prop if `passNode` is on).
  *
  * You can access props at `JSX.IntrinsicElements`.
  * For example, to find props for `a`, use `JSX.IntrinsicElements['a']`.
@@ -53,6 +57,6 @@ export type Component<ComponentProps> =
 // react into the `.d.ts` file.
 export type Components = {
   [TagName in keyof JSX.IntrinsicElements]:
-    | Component<JSX.IntrinsicElements[TagName]>
+    | Component<JSX.IntrinsicElements[TagName] & ExtraProps>
     | keyof JSX.IntrinsicElements
 }
diff --git a/lib/index.js b/lib/index.js
index 6cee227..5f7e143 100644
--- a/lib/index.js
+++ b/lib/index.js
@@ -18,7 +18,7 @@
  * @callback Jsx
  *   Create a production element.
  * @param {unknown} type
- *   Element type: the `Fragment` symbol or a tag name (`string`).
+ *   Element type: `Fragment` symbol, tag name (`string`), component.
  * @param {Props} props
  *   Element props and also includes `children`.
  * @param {string | undefined} key
@@ -29,7 +29,7 @@
  * @callback JsxDev
  *   Create a development element.
  * @param {unknown} type
- *   Element type: the `Fragment` symbol or a tag name (`string`).
+ *   Element type: `Fragment` symbol, tag name (`string`), component.
  * @param {Props} props
  *   Element props and also includes `children`.
  * @param {string | undefined} key
@@ -67,7 +67,7 @@
  * @typedef {JSX.Element | string | null | undefined} Child
  *   Child.
  *
- * @typedef {{children: Array<Child>, [prop: string]: Value | Array<Child>}} Props
+ * @typedef {{children: Array<Child>, node?: Element | undefined, [prop: string]: Value | Element | undefined | Array<Child>}} Props
  *   Properties and children.
  *
  * @callback Create
@@ -89,6 +89,8 @@
  *   File path.
  * @property {Partial<Components>} components
  *   Components to swap.
+ * @property {boolean} passNode
+ *   Pass `node` to components.
  * @property {Schema} schema
  *   Current schema.
  * @property {unknown} Fragment
@@ -108,6 +110,8 @@
  *
  *   Passed in source info to `jsxDEV` when using the automatic runtime with
  *   `development: true`.
+ * @property {boolean | null | undefined} [passNode=false]
+ *   Pass the hast element node to components.
  * @property {Space | null | undefined} [space='html']
  *   Whether `tree` is in the `'html'` or `'svg'` space.
  *
@@ -239,6 +243,7 @@ export function toJsxRuntime(tree, options) {
   const state = {
     Fragment: options.Fragment,
     schema: options.space === 'svg' ? svg : html,
+    passNode: options.passNode || false,
     components: options.components || {},
     filePath,
     create
@@ -298,6 +303,11 @@ function one(state, node, key) {
       if (own.call(state.components, node.tagName)) {
         const key = /** @type {keyof JSX.IntrinsicElements} */ (node.tagName)
         type = state.components[key]
+
+        // If this is swapped out for a component:
+        if (type !== 'string' && state.passNode) {
+          props.node = node
+        }
       } else {
         type = node.tagName
       }
diff --git a/readme.md b/readme.md
index d57d9c9..7608248 100644
--- a/readme.md
+++ b/readme.md
@@ -144,6 +144,11 @@ Static JSX ([`Jsx`][jsx], required in production).
 
 Development JSX ([`JsxDev`][jsxdev], required in development).
 
+###### `development`
+
+Whether to use `jsxDEV` when on or `jsx` and `jsxs` when off (`boolean`,
+default: `false`).
+
 ###### `components`
 
 Components to use ([`Partial<Components>`][components], optional).
@@ -151,11 +156,6 @@ Components to use ([`Partial<Components>`][components], optional).
 Each key is the name of an HTML (or SVG) element to override.
 The value is the component to render instead.
 
-###### `development`
-
-Whether to use `jsxDEV` when on or `jsx` and `jsxs` when off (`boolean`,
-default: `false`).
-
 ###### `filePath`
 
 File path to the original source file (`string`, optional).
@@ -163,6 +163,10 @@ File path to the original source file (`string`, optional).
 Passed in source info to `jsxDEV` when using the automatic runtime with
 `development: true`.
 
+###### `passNode`
+
+Pass the hast element node to components (`boolean`, default: `false`).
+
 ###### `space`
 
 Whether `tree` is in the `'html'` or `'svg'` space ([`Space`][space], default:
@@ -183,8 +187,8 @@ it.
 Possible components to use (TypeScript type).
 
 Each key is a tag name typed in `JSX.IntrinsicElements`.
-Each value is a component accepting the corresponding props or a different tag
-name.
+Each value is either a different tag name, or a component accepting the
+corresponding props (and an optional `node` prop if `passNode` is on).
 
 You can access props at `JSX.IntrinsicElements`.
 For example, to find props for `a`, use `JSX.IntrinsicElements['a']`.
@@ -192,12 +196,16 @@ For example, to find props for `a`, use `JSX.IntrinsicElements['a']`.
 ###### Type
 
 ```ts
+import type {Element} from 'hast'
+
 type Components = {
   [TagName in keyof JSX.IntrinsicElements]:
-    | Component<JSX.IntrinsicElements[TagName]>
+    | Component<JSX.IntrinsicElements[TagName] & ExtraProps>
     | keyof JSX.IntrinsicElements
 }
 
+type ExtraProps = {node?: Element | undefined}
+
 type Component<ComponentProps> =
   // Function component:
   | ((props: ComponentProps) => JSX.Element | string | null | undefined)
@@ -222,7 +230,7 @@ Create a production element (TypeScript type).
 ###### Parameters
 
 *   `type` (`unknown`)
-    — element type: the `Fragment` symbol or a tag name (`string`)
+    — element type: `Fragment` symbol, tag name (`string`), component
 *   `props` ([`Props`][props])
     — element props and also includes `children`
 *   `key` (`string` or `undefined`)
@@ -239,7 +247,7 @@ Create a development element (TypeScript type).
 ###### Parameters
 
 *   `type` (`unknown`)
-    — element type: the `Fragment` symbol or a tag name (`string`)
+    — element type: `Fragment` symbol, tag name (`string`), component
 *   `props` ([`Props`][props])
     — element props and also includes `children`
 *   `key` (`string` or `undefined`)
diff --git a/test/index.js b/test/index.js
index cd0d728..e98a7cc 100644
--- a/test/index.js
+++ b/test/index.js
@@ -400,6 +400,39 @@ test('components', () => {
     'a',
     'should support class components'
   )
+
+  assert.equal(
+    renderToStaticMarkup(
+      toJsxRuntime(h('b'), {
+        ...production,
+        passNode: true,
+        components: {
+          b(props) {
+            assert.ok(props.node)
+            return 'a'
+          }
+        }
+      })
+    ),
+    'a',
+    'should support components w/ `passNode: true`'
+  )
+
+  assert.equal(
+    renderToStaticMarkup(
+      toJsxRuntime(h('b'), {
+        ...production,
+        components: {
+          b(props) {
+            assert.equal(props.node, undefined)
+            return 'a'
+          }
+        }
+      })
+    ),
+    'a',
+    'should support components w/o `passNode`'
+  )
 })
 
 test('react specific: filter whitespace in tables', () => {