TS infer wrong type for function

Question

Hello everyone i have a code like this:

import { z } from "zod"

export const { object, string, array, union, literal } = z

type AddressKind = "residence" | "legal"

export function address<TKind extends AddressKind>(kind: TKind) {
    const base = object({
        street: z.string().max(50),
    })

    switch (kind) {
        case "residence": {
            return base.merge(
                object({
                    flat: z.optional(z.string().max(10)),
                }),
            )
        }

        case "legal": {
            return base.merge(
                object({
                    office: z.optional(z.string().max(10)),
                }),
            )
        }

        default: {
            throw new Error("Unexpected kind")
        }
    }
}

export type Address<Kind extends AddressKind> = z.infer<ReturnType<typeof address<Kind>>>

export type LegalAdress = Address<"legal">
// type LegalAdress = {
//     street: string;
//     flat?: string | undefined;
// } | {
//     street: string;
//     office?: string | undefined;
// }

I expect that LegalAdress type is

{
    street: string;
    office?: string | undefined;
}

but right now it is a union:

type LegalAdress = {
 street: string;
 flat?: string | undefined;
} | {
 street: string;
 office?: string | undefined;
}

why does this happen and how to fix it?

Playground

Answer 1

You can streamline the code by creating a mapping type that links each AddressKind to its respective schema. This mapping object can also serve as the source for deriving the possible values of the AddressKind type.

This should work:

Note: The throw new Error line does not even need to be in address, because you will get the type at compile time. You can simply return schemas[kind].

import { z } from "zod";

const baseSchema = z.object({
    street: z.string().max(50),
});

const schemas = {
    residence: baseSchema.merge(
        z.object({
            flat: z.optional(z.string().max(10)),
        })
    ),
    legal: baseSchema.merge(
        z.object({
            office: z.optional(z.string().max(10)),
        })
    ),
} as const;

export type AddressKind = keyof typeof schemas; // "residence" | "legal"

export function address(kind: AddressKind) {
    const schema = schemas[kind];
    if (!schema) {
        throw new Error("Unexpected kind"); // This is not really needed
    }
    return schema;
}

export type Address<Kind extends AddressKind> = z.infer<typeof schemas[Kind]>;

export type LegalAddress = Address<"legal">;
export type ResidenceAddress = Address<"residence">;

Inferred types:

type AddressKind = "residence" | "legal"

type LegalAddress = {
    street: string;
    office?: string | undefined;
}

type ResidenceAddress = {
    street: string;
    flat?: string | undefined;
}

Usage:

const legalAddressData: LegalAddress = {
    street: "123 Legal St",
    office: "Suite 100",
};

const residenceAddressData: ResidenceAddress = {
    street: "456 Residence Ave",
    flat: "Apt 2B",
};

// Validate the data using the Zod schemas
const legalAddress = address('legal').parse(legalAddressData);
const residenceAddress = address('residence').parse(residenceAddressData);