TypeScript函数返回类型推断错误问题排查与修复

问题描述 #

编写了如下Zod代码，期望LegalAdress类型为仅包含street和可选office的对象，但实际得到的是residence和legal地址类型的联合类型：

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;
// }

期望的LegalAdress类型：

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

问题原因 #

TypeScript无法根据泛型参数TKind自动缩小address函数的返回类型。虽然函数内部通过switch case根据kind参数返回不同的Zod schema，但编译器无法将泛型类型与具体返回值类型建立关联，因此推断出的返回类型是两种schema的联合类型，最终导致z.infer得到的也是联合类型。

修复方案 #

以下两种方式均可解决问题：

方式一：使用函数重载 #

通过函数重载明确指定每种kind对应的返回类型，让TypeScript准确推断：

import { z } from "zod"

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

type AddressKind = "residence" | "legal"

// 函数重载定义
export function address(kind: "residence"): z.ZodObject<{
    street: z.ZodString;
    flat: z.ZodOptional<z.ZodString>;
}>;
export function address(kind: "legal"): z.ZodObject<{
    street: z.ZodString;
    office: z.ZodOptional<z.ZodString>;
}>;
// 实现函数
export function address(kind: AddressKind) {
    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">
// 现在类型正确：{ street: string; office?: string | undefined; }

方式二：使用条件类型约束返回类型 #

通过条件类型映射TKind到对应的schema类型，给函数添加明确的返回类型：

import { z } from "zod"

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

type AddressKind = "residence" | "legal"

// 定义每种类型对应的schema类型
type AddressSchema<T extends AddressKind> = 
    T extends "residence" ? z.ZodObject<{
        street: z.ZodString;
        flat: z.ZodOptional<z.ZodString>;
    }> :
    T extends "legal" ? z.ZodObject<{
        street: z.ZodString;
        office: z.ZodOptional<z.ZodString>;
    }> : never;

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

    switch (kind) {
        case "residence": {
            return base.merge(
                object({
                    flat: z.optional(z.string().max(10)),
                }),
            ) as AddressSchema<TKind>;
        }

        case "legal": {
            return base.merge(
                object({
                    office: z.optional(z.string().max(10)),
                }),
            ) as AddressSchema<TKind>;
        }

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

export type Address<Kind extends AddressKind> = z.infer<AddressSchema<Kind>>

export type LegalAdress = Address<"legal">
// 现在类型正确：{ street: string; office?: string | undefined; }

内容的提问来源于stack exchange，提问作者Pavel Perevezencev