Infer return type of function that takes a discriminated union?

Question

I have a function that takes in an object that uses a discriminated union that returns a different structure based on that property.

While i could write an assertion isLicenseType fn, I was wondering if there is a way I could infer this based on the type argument provided to the fn, possibly with an assertion signature?

const licenseTypeMap: Record<string, LicenseOptions> = { abc: { ... } }
const siteIdMap: Record<number, SiteIdOptions> = { 12: { ... } }

type DynamicReferenceType =
    | { type: 'LICENSE'; licenseName?: string; }
    | { type: 'SITE_ID'; siteId?: number; };

export function getDynamicSchema(args: DynamicReferenceType) {
    if (args.type === 'LICENSE') {
        return licenseTypeMap[args.licenseName];
    }
    return siteIdMap[args.siteId];
}

const schema = getDynamicSchema({ type: 'license', licenseName: 'abc' })
// schema is of LicenseOptions | SiteIdOptions

Answer 1

Came across a solution that's as close as I'm going to get it seems. TS can't narrow the type of generic parameters, so there's a little bit of a casting required but:

const licenseTypeMap: Record<string, LicenseOptions> = { abc: { ... } }
const siteIdMap: Record<number, SiteIdOptions> = { 12: { ... } }

type DynamicReferenceType =
    | { type: 'LICENSE'; licenseName?: string; }
    | { type: 'SITE_ID'; siteId?: number; };

type ReturnTypeLookup = {
    LICENSE: LicenseOptions;
    SITE_ID: SiteIdOptions;
};

type SchemaFor<T extends DynamicReferenceType> = ReturnTypeLookup[T['type']];

export function getDynamicComponent<T extends DynamicReferenceType>(args: T): SchemaFor<T> {
    if (args.type === 'LICENSE') {
        return licenseTypeMap[args.licenseName] as SchemaFor<T>;
    }
    return siteIdMap[args.siteId] as SchemaFor<T>;
}