Some zod schema utilities for creating a Zod schema for a recursive structure

base-node.ts

import { z } from "zod";

export type BaseNode<T extends string = "children"> = {
  [K in T]?: BaseNode<T>[];
};

/**
 * Creates a schema for a node that forms a directed tree.
 *
 * Usage:
 *
 *     const SimpleTreeNode = baseNode("children", m => m))
 *
 *     type SimpleTreeNode = z.infer<typeof SimpleTreeNode>
 *
 * If you want the children to be either a node or something else (such as a string):
 *
 *     const NodeOrString = baseNode("children", m => z.union([m, z.string()]));
 *
 *     type NodeOrString = z.infer<typeof NodeOrString>
 *
 * If you want to add fields to each node:
 *
 *     const NodeWithStyle = baseNode("children", m => m.and(z.object({ style: z.unknown() })))
 *
 *     type NodeWithStyle = z.infer<typeof NodeWithStyle>
 * @param childrenKey A string that represents the field name that will be a list of children.
 * @param modifier Useful for adding fields to each child node
 * @returns A zod schema
 */
export const baseNode = <R, V extends string = "children">(
  childrenKey: V,
  modifier: (k: z.ZodType<BaseNode<V>>) => z.ZodType<R>
) =>
  modifier(
    z.lazy(
      (): z.ZodType<BaseNode<V>> =>
        z.object({
          [childrenKey]: z.array(baseNode(childrenKey, modifier)).optional(),
        }) as unknown as z.ZodType<BaseNode<V>>
    )
  );