TypeScript: Generic function return type inference issue with zod schema validation

Question

I'm using the zod library for schema validation and have defined two schemas: FormFactorDevice and FormFactor. I've also implemented a selectFormFactorDevice function to select a value from the FormFactorDevice object based on a given deviceType.

However, the return value of selectFormFactorDevice is inferred as any type, whereas I expect it to be string | undefined since I'm using z.string() as the fieldSchema parameter.

Here's the code:

export const FormFactorDevice = <T extends z.ZodTypeAny>(fieldSchema: T) =>
  z.object({
    all: z.optional(fieldSchema),
    phone: z.optional(fieldSchema),
    smalltablet: z.optional(fieldSchema),
    tablet: z.optional(fieldSchema),
    laptop: z.optional(fieldSchema),
    desktop: z.optional(fieldSchema),
  });

export const FormFactor = <Z extends z.ZodTypeAny>(fieldSchema: Z) =>
  z.object({
    isFormFactor: z.literal(true).default(true),
    all: z.optional(fieldSchema),
    landscape: z.optional(FormFactorDevice(fieldSchema)),
    portrait: z.optional(FormFactorDevice(fieldSchema)),
  });

// FormFactorDevice could receive very complicated schema
const b = FormFactorDevice(z.string()).parse({
  all: '100',
  smalltablet: '200',
  phone: '300',
});

export const selectFormFactorDevice = <T extends z.ZodTypeAny>(
  formFactorDevice: z.infer<ReturnType<typeof FormFactorDevice<T>>>,
  deviceType: 'all' | 'phone' | 'smalltablet' | 'tablet' | 'laptop' | 'desktop'
): z.infer<T> | undefined => {
  return formFactorDevice[deviceType];
};

const d = selectFormFactorDevice(b, 'all'); // d is inferred as `any` type

How can I modify the selectFormFactorDevice function to return the correct type (string | undefined in this case)?

Note: I've tried using z.infer and ReturnType<typeof FormFactorDevice> to infer the correct type, but it doesn't seem to work. The return value of selectFormFactorDevice is inferred as any type, whereas I expect it to be string | undefined since I'm using z.string() as the fieldSchema parameter.

I know selectFormFactorDevice<z.ZodString>(b, 'all') should work but I want my function to be smart enough to infer the type automatically

You can try copy and paste the code to zod playground here: https://stackblitz.com/edit/typescript-rqwgxo?file=index.ts

Answer 1

To ensure that the selectFormFactorDevice function correctly infers the return type as string | undefined, you need to adjust the generic type inference. Currently, TypeScript struggles to infer the correct type due to the complexity of your generic type parameter T.

You can achieve the desired inference by making slight adjustments to your function signature. Instead of directly passing the fieldSchema parameter to FormFactorDevice, you can explicitly specify the type for each form factor in FormFactorDevice. Here's how you can modify your code:

import * as z from 'zod';

export const FormFactorDevice = <T extends z.ZodTypeAny>(fieldSchema: T) =>
  z.object({
    all: z.optional(fieldSchema),
    phone: z.optional(fieldSchema),
    smalltablet: z.optional(fieldSchema),
    tablet: z.optional(fieldSchema),
    laptop: z.optional(fieldSchema),
    desktop: z.optional(fieldSchema),
  });

export const FormFactor = <Z extends z.ZodTypeAny>(fieldSchema: Z) =>
  z.object({
    isFormFactor: z.literal(true).default(true),
    all: z.optional(fieldSchema),
    landscape: z.optional(FormFactorDevice(z.object({ landscape: fieldSchema }))),
    portrait: z.optional(FormFactorDevice(z.object({ portrait: fieldSchema }))),
  });

// Adjusted the return type of selectFormFactorDevice
export const selectFormFactorDevice = <T extends z.ZodTypeAny>(
  formFactorDevice: z.infer<ReturnType<typeof FormFactorDevice<T>>>,
  deviceType: keyof typeof formFactorDevice
): z.infer<T> | undefined => {
  return formFactorDevice[deviceType];
};

// Example usage:
const b = FormFactorDevice(z.string()).parse({
  all: '100',
  smalltablet: '200',
  phone: '300',
});

const d = selectFormFactorDevice(b, 'all'); // d is inferred as `string | undefined` type

By explicitly specifying the type for each form factor in FormFactorDevice, TypeScript can better infer the return type of selectFormFactorDevice. Now, the selectFormFactorDevice function correctly infers the return type as string | undefined based on the provided fieldSchema.

Answer 2

Note: This answer assumes that the TS strictNullChecks compiler option is set to true (which is not the case for your StackBlitz link).

The selectFormFactorDevice function has a very simple implementation (it is simply a property access). I assume you want the typings for it to do 2 things:

1. Infer the correct return type.
2. Restrict the first argument to objects having been parsed by a schema created using the FormFactorDevice function.

I think these typings will fulfill both of those requirements at least somewhat well:

export const selectFormFactorDevice = <
  TFormFactorDevice extends z.infer<
    ReturnType<typeof FormFactorDevice<ZodTypeAny>>
  >,
  TDeviceType extends
    | 'all'
    | 'phone'
    | 'smalltablet'
    | 'tablet'
    | 'laptop'
    | 'desktop'
>(
  formFactorDevice: TFormFactorDevice,
  deviceType: TDeviceType
) => {
  return formFactorDevice[deviceType];
};

The constraint on TFormFactorDevice expresses point 1 (although it may do so in an incomplete manner; I am no expert in the Zod typings). Using a generic parameter TDeviceType for deviceType ensures that the return type is correctly inferred as well (point 2); it simply becomes TFormFactorDevice[TDeviceType].

If you don't want to type out all the string options for TDeviceType, you can also use the keyof type operator like this: TDeviceType extends keyof TFormFactorDevice.