108 lines
3 KiB
TypeScript
108 lines
3 KiB
TypeScript
import { jPatches } from './_internal'
|
|
import type { ComponentProps, ElementType, Key, ReactElement } from 'react'
|
|
import type { ReactJSXRuntime } from '..'
|
|
|
|
export type AnyJSXFactoryFunction = (typeof ReactJSXRuntime)['jsx' | 'jsxs']
|
|
|
|
export type BeforeJSXCallback<
|
|
E extends ElementType = ElementType,
|
|
P = ComponentProps<E>,
|
|
> = (
|
|
args: [element: E, props: P, key?: Key | undefined],
|
|
) => Parameters<AnyJSXFactoryFunction>
|
|
|
|
export type InsteadJSXCallback<
|
|
E extends ElementType = ElementType,
|
|
P = ComponentProps<E>,
|
|
> = (
|
|
args: [element: E, props: P, key?: Key | undefined],
|
|
jsx: AnyJSXFactoryFunction,
|
|
) => ReturnType<AnyJSXFactoryFunction> | null
|
|
|
|
export type AfterJSXCallback = (
|
|
fiber: ReactElement,
|
|
) => ReturnType<AnyJSXFactoryFunction> | null
|
|
|
|
/**
|
|
* Registers a hook to be called after a JSX element with the specified type is created.
|
|
*
|
|
* @param type The type of the element.
|
|
* @param patch The hook.
|
|
* @returns A function to unpatch.
|
|
*/
|
|
export function afterJSX(type: ElementType, patch: AfterJSXCallback) {
|
|
let patches = jPatches.get(type)
|
|
if (!patches) {
|
|
patches = [undefined, new Set(), undefined]
|
|
jPatches.set(type, patches)
|
|
}
|
|
|
|
const set = patches[1]!
|
|
|
|
set.add(patch as AfterJSXCallback)
|
|
return () => {
|
|
const res = set.delete(patch as AfterJSXCallback)
|
|
if (res) attemptCleanup(type)
|
|
return res
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Registers a hook to be called before a JSX element with the specified type is created.
|
|
*
|
|
* @param type The type of the element.
|
|
* @param patch The hook.
|
|
* @returns A function to unpatch.
|
|
*/
|
|
export function beforeJSX<
|
|
E extends ElementType = ElementType,
|
|
P = ComponentProps<E>,
|
|
>(type: E, patch: BeforeJSXCallback<E, P>) {
|
|
let patches = jPatches.get(type)
|
|
if (!patches) {
|
|
patches = [new Set(), undefined, undefined]
|
|
jPatches.set(type, patches)
|
|
}
|
|
|
|
const set = patches[0]!
|
|
|
|
set.add(patch as BeforeJSXCallback)
|
|
return () => {
|
|
const res = set.delete(patch as BeforeJSXCallback)
|
|
if (res) attemptCleanup(type)
|
|
return res
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Registers a callback to run instead when a JSX element with the specified type is created.
|
|
*
|
|
* @param type The type of the element.
|
|
* @param patch The hook.
|
|
* @returns A function to unpatch.
|
|
*/
|
|
export function insteadJSX<
|
|
E extends ElementType = ElementType,
|
|
P = ComponentProps<E>,
|
|
>(type: E, patch: InsteadJSXCallback<E, P>) {
|
|
let patches = jPatches.get(type)
|
|
if (!patches) {
|
|
patches = [undefined, undefined, new Set()]
|
|
jPatches.set(type, patches)
|
|
}
|
|
|
|
const set = patches[2]!
|
|
|
|
set.add(patch as InsteadJSXCallback)
|
|
return () => {
|
|
const res = set.delete(patch as InsteadJSXCallback)
|
|
if (res) attemptCleanup(type)
|
|
return res
|
|
}
|
|
}
|
|
|
|
function attemptCleanup(type: ElementType) {
|
|
const patches = jPatches.get(type)!
|
|
if (patches[0]?.size || patches[1]?.size || patches[2]?.size) return
|
|
jPatches.delete(type)
|
|
}
|