How to narrow down the type of register (UseFormRegister) within a form which uses a discriminated union as its FieldValues.

[virtuallyunknown]

Greetings!

I am trying to create a dynamic form which accepts this type as useForm generic param:

type AnimalSchema = {
    type: "cat";
    meow: string;
} | {
    type: "dog";
    bark: string;
} | {
    type: "bird";
    chirp: string;
}

The idea here is very simple. Have a normal <select> element which changes the value of type, then display different inputs depending on that type. If it's a cat, we want to show <input {...register('cat') />, if it's a dog, we want <input {...register('dog') />, you get the idea.

As far as runtime behavior is concerned, this all works, my question relates to Typescript exclusively. I want to be able to narrow down the possible type of register.

export default function App() {
  const { register, watch } = useForm<AnimalSchema>({
    defaultValues: { type: "bird" },
    mode: "onChange",
  });

  const animalType = watch("type");

  return (
    <form>
      {formData.type === "bird" ? (
          <input type="text" {...register("bark")} />
      ) : null}
    </form>
  );
}

In the example above, we are attempting to narrow down register, however, it's still being inferred by typescript as a union of:

const register: <"bark">(name: "bark", options?: RegisterOptions<{
    type: "cat";
    meow: string;
} | {
    type: "dog";
    bark: string;
} | {
    type: "bird";
    chirp: string;
}, "bark"> | undefined) => UseFormRegisterReturn<...>

and thus, I can {...register("bark")}, even though bark is not a valid property for this branch of the union. This is exactly what I want to prevent, in order to enforce more type safety.

I made a codesandbox demo, and I would appreciate it if someone has any info about this.

Cheers!

[virtuallyunknown]

I finally managed to find a semi-decent solution, and I am posting it here for future reference and in case someone else finds this. But I am honestly kind of shocked something as common as this requires so much time digging through Discord threads, online forums and issues but no viable solution to be found. Do people only build linear forms these days, or maybe I am missing something obvious. And if you think tanstack form solves this issue, think again.

import "./styles.css";
import { useForm, UseFormReturn } from "react-hook-form";
import { z } from "zod";
import { InputSelect } from "./Select";
import { InputText } from "./Text";
import { zodResolver } from "@hookform/resolvers/zod";

const birdSchema = z.object({
  type: z.literal("bird"),
  chirp: z.string().min(5),
});
const catSchema = z.object({ type: z.literal("cat"), meow: z.string().min(5) });
const dogSchema = z.object({ type: z.literal("dog"), bark: z.string().min(5) });

const animalSchema = z.discriminatedUnion("type", [
  birdSchema,
  catSchema,
  dogSchema,
]);

type AnimalSchema = z.infer<typeof animalSchema>;
type BirdFormBranch = UseFormReturn<Extract<AnimalSchema, { type: "bird" }>>;
type CatFormBranch = UseFormReturn<Extract<AnimalSchema, { type: "cat" }>>;
type DogFormBranch = UseFormReturn<Extract<AnimalSchema, { type: "dog" }>>;

const BirdForm = <T extends BirdFormBranch>({ form }: { form: T }) => {
  return (
    <label>
      <span>Bird Form</span>
      <InputText name="chirp" register={form.register} type="text" />
      {form.formState.errors?.chirp ? (
        <span>{form.formState.errors?.chirp.message}</span>
      ) : null}
    </label>
  );
};

const CatForm = <T extends CatFormBranch>({ form }: { form: T }) => {
  return (
    <label>
      <span>Cat Form</span>
      <input {...form.register("meow")} />
      {form.formState.errors?.meow ? (
        <span>{form.formState.errors?.meow.message}</span>
      ) : null}
    </label>
  );
};

const DogForm = <T extends DogFormBranch>({ form }: { form: T }) => {
  return (
    <label>
      <span>Dog Form</span>
      <input {...form.register("bark")} />
      {form.formState.errors?.bark ? (
        <span>{form.formState.errors?.bark.message}</span>
      ) : null}
    </label>
  );
};

export default function App() {
  const form = useForm<AnimalSchema>({
    defaultValues: { type: "bird" },
    mode: "onChange",
    resolver: zodResolver(animalSchema),
  });

  const { register, watch } = form;

  const animalType = watch("type");

  const subForm = (() => {
    switch (animalType) {
      case "cat":
        return <CatForm form={form as CatFormBranch} />;
      case "dog":
        return <DogForm form={form as DogFormBranch} />;
      case "bird":
        return <BirdForm form={form as BirdFormBranch} />;
      default:
        throw new Error("Invalid form type");
    }
  })();

  return (
    <form>
      <h1>Animal form</h1>
      <label>
        <span>Select animal type</span>
        <InputSelect
          name="type"
          register={register}
          options={[
            { text: "Bird", value: "bird" },
            { text: "Cat", value: "cat" },
            { text: "Dog", value: "dog" },
          ]}
        />
      </label>
      {subForm}
    </form>
  );
}

You first need to create types for all the branches of the discriminated union - BirdFormBranch, CatFormBranch, DogFormBranch. Then you create subforms components for each of one of them. The key part here is that once you are in the component itself, for example CatForm, then type safety is guaranteed. You can't use an input <input {...form.register("meow")} />, TS will yell at you. And you also get the formState.errors in their correct state.

Inside the main component, you do an IIFE switch for convenience. Unfortunately type assertion here is unavoidable, but I think this is the only "weak-link", at least as far as I've discovered.

Here is an working demo, it works with zod validation as well:

https://codesandbox.io/p/sandbox/qksgsj

All the best.