feat: support ZodUnions of ZodLiterals (#7)

Co-authored-by: Will Ruggiano <wmruggiano@gmail.com>

packages/dev-app/src/router.ts

@@ -90,6 +90,15 @@ export const appRouter = createTRPCRouter({
       .query(() => {
         return "It's an input";
       }),
+    unionInput: procedure
+      .input(
+        z.object({
+          aUnion: z.union([z.literal("one"), z.literal(2)]),
+        })
+      )
+      .query(({ input }) => {
+        return input;
+      }),
     emailTextInput: procedure
       .input(
         z.object({
@@ -175,7 +184,7 @@ export const appRouter = createTRPCRouter({
         optionalEnum: z.enum(["Three", "Four"]).optional(),
         stringArray: z.string().array(),
         boolean: z.boolean(),
-        union: z.discriminatedUnion("disc", [
+        discriminatedUnion: z.discriminatedUnion("disc", [
           z.object({
             disc: z.literal("one"),
             oneProp: z.string(),
@@ -185,6 +194,7 @@ export const appRouter = createTRPCRouter({
             twoProp: z.enum(["one", "two"]),
           }),
         ]),
+        union: z.union([z.literal("one"), z.literal(2)]),
       })
     )
     .query(() => ({ goodJob: "yougotthedata" })),

packages/trpc-panel/src/parse/input-mappers/__tests__/zod/union.test.ts

@@ -0,0 +1,28 @@
+import { defaultReferences } from "@src/parse/input-mappers/defaultReferences";
+import { parseZodUnionDef } from "@src/parse/input-mappers/zod/parsers/parseZodUnionDef";
+import { UnionNode } from "@src/parse/parseNodeTypes";
+import { z } from "zod";
+
+describe("Parse Zod Union", () => {
+  it("should parse a union node", () => {
+    const expected: UnionNode = {
+      type: "union",
+      path: [],
+      values: [
+        {
+          type: "literal",
+          value: "one",
+          path: [],
+        },
+        {
+          type: "literal",
+          value: 2,
+          path: [],
+        },
+      ],
+    };
+    const zodSchema = z.union([z.literal("one"), z.literal(2)]);
+    const parsedZod = parseZodUnionDef(zodSchema._def, defaultReferences());
+    expect(parsedZod).toStrictEqual(expected);
+  });
+});

packages/trpc-panel/src/parse/input-mappers/zod/parsers/parseZodUnionDef.ts

@@ -0,0 +1,18 @@
+import { nodePropertiesFromRef } from "@src/parse/utils";
+import { ZodUnionDef } from "zod";
+import { UnionNode, ParseFunction, LiteralNode } from "../../../parseNodeTypes";
+import { zodSelectorFunction } from "../selector";
+
+export const parseZodUnionDef: ParseFunction<ZodUnionDef, UnionNode> = (
+  def,
+  refs
+) => {
+  refs.addDataFunctions.addDescriptionIfExists(def, refs);
+  return {
+    type: "union",
+    values: def.options.map(
+      (o) => zodSelectorFunction(o._def, { ...refs, path: [] }) as LiteralNode
+    ),
+    ...nodePropertiesFromRef(refs),
+  };
+};

packages/trpc-panel/src/parse/input-mappers/zod/selector.ts

@@ -16,6 +16,7 @@ import {
   ZodPromiseDef,
   ZodStringDef,
   ZodUndefinedDef,
+  ZodUnionDef,
   ZodVoidDef,
 } from "zod";
 import { parseZodStringDef } from "./parsers/parseZodStringDef";
@@ -40,6 +41,7 @@ import { parseZodEffectsDef } from "@src/parse/input-mappers/zod/parsers/parseZo
 import { parseZodNullDef } from "@src/parse/input-mappers/zod/parsers/parseZodNullDef";
 import { parseZodPromiseDef } from "@src/parse/input-mappers/zod/parsers/parseZodPromiseDef";
 import { parseZodUndefinedDef } from "@src/parse/input-mappers/zod/parsers/parseZodUndefinedDef";
+import { parseZodUnionDef } from "@src/parse/input-mappers/zod/parsers/parseZodUnionDef";
 import { parseZodVoidDef } from "./parsers/parseZodVoidDef";
 
 export const zodSelectorFunction: ParserSelectorFunction<ZodDefWithType> = (
@@ -89,6 +91,8 @@ export const zodSelectorFunction: ParserSelectorFunction<ZodDefWithType> = (
       return parseZodPromiseDef(def as ZodPromiseDef, references);
     case ZodFirstPartyTypeKind.ZodUndefined:
       return parseZodUndefinedDef(def as ZodUndefinedDef, references);
+    case ZodFirstPartyTypeKind.ZodUnion:
+      return parseZodUnionDef(def as ZodUnionDef, references);
     case ZodFirstPartyTypeKind.ZodVoid:
       return parseZodVoidDef(def as ZodVoidDef, references);
   }

packages/trpc-panel/src/parse/parseNodeTypes.ts

@@ -33,6 +33,11 @@ export type DiscriminatedUnionNode = {
   discriminatorName: string;
 } & SharedInputNodeProperties;
 
+export type UnionNode = {
+  type: "union";
+  values: LiteralNode[];
+} & SharedInputNodeProperties;
+
 /**
  * Any time you just want the front end to send back a value use this
  */
@@ -58,6 +63,7 @@ export type ParsedInputNode =
   | ObjectNode
   | EnumNode
   | DiscriminatedUnionNode
+  | UnionNode
   | LiteralNode
   | StringNode
   | NumberNode

packages/trpc-panel/src/react-app/components/form/Field.tsx

@@ -10,6 +10,7 @@ import { LiteralField } from "./fields/LiteralField";
 import { NumberField } from "./fields/NumberField";
 import { ObjectField } from "./fields/ObjectField";
 import { TextField } from "./fields/TextField";
+import { UnionField } from "./fields/UnionField";
 
 export function Field({
   inputNode,
@@ -77,6 +78,15 @@ export function Field({
           node={inputNode}
         />
       );
+    case "union":
+      return (
+        <UnionField
+          name={path}
+          label={label}
+          control={control}
+          node={inputNode}
+        />
+      );
     case "literal":
       return <LiteralField />;
     case "unsupported":

packages/trpc-panel/src/react-app/components/form/fields/UnionField.tsx

@@ -0,0 +1,31 @@
+import React from "react";
+import { Control, useController } from "react-hook-form";
+import type { ParsedInputNode } from "@src/parse/parseNodeTypes";
+import { BaseSelectField } from "./base/BaseSelectField";
+
+export function UnionField({
+  name,
+  label,
+  control,
+  node,
+}: {
+  name: string;
+  label: string;
+  control: Control<any>;
+  node: ParsedInputNode & { type: "union" };
+}) {
+  const { field, fieldState } = useController({
+    name,
+    control,
+  });
+
+  return (
+    <BaseSelectField
+      options={node.values.map((n) => n.value as string)}
+      value={field.value}
+      onChange={field.onChange}
+      errorMessage={fieldState.error?.message}
+      label={label}
+    />
+  );
+}