Infer type of argument and return type based on other argument

Question

A function takes an object which can be either of TypeA or TypeB. The first parameter specifies the type of the object. The returned type also depends on this first parameter.

Problem: TypeScript does not infer the type of the object within a case (or within an if). The following code has errors.

I know I can fix this issue by be explicit about the types (add as TypeA ) but I wanted to know if there are other solutions to avoid having to put the explicit cast.

type TypeA = {
  input: { foo: string }
  output: { foo: number }
}

type TypeB = {
  input: { bar: string }
  output: { bar: number }
}

type Types = {
  TypeA: TypeA
  TypeB: TypeB
}

type TypesName = keyof Types

/**
 * Takes TypeA or TypeB objects, and the return type is based on the type.
 */
function transform<N extends TypesName>(typeName: N, typeValue: Types[N]['input']): Types[N]['output'] {
  switch (typeName) {
    case 'TypeA':
      return transformTypeA(typeValue) // Argument of type '{ foo: string; } | { bar: string; }' is not assignable to parameter of type '{ foo: string; }'.
    case 'TypeB':
      return transformTypeB(typeValue) // Argument of type '{ foo: string; } | { bar: string; }' is not assignable to parameter of type '{ bar: string; }'.
  }
  throw new Error('Unknown type')
}

function transformTypeA(typeA: TypeA['input']): TypeA['output'] {
  return { foo: parseInt(typeA.foo) }
}

function transformTypeB(typeB: TypeB['input']): TypeB['output'] {
  return { bar: parseInt(typeB.bar) }
}


const transformedValue = transform('TypeA', { foo: 'lol' })
console.log(transformedValue) // transformedValue is of type { foo: number }

Answer 1

There is a way but it's probably not worth right now.

type Union<N extends keyof Types = keyof Types> = N extends N ? [typeName: N, typeValue: Types[N]['input']]: never;

function transform<N extends TypesName>(...p: Union<N>): Types[N]['output'] {
  switch (p[0]) {
    case 'TypeA':
      return transformTypeA(p[1])
    case 'TypeB':
      return transformTypeB(p[1])
  }
}

Playground Link

We need to use index access instead of defining separate parameters or destructuring the parameters as this is the only way typescript will see the two values are related. This might be supported at a later date as described here here

In particular, the pattern of destructuring a discriminant property and a payload property into two local variables and expecting a coupling between the two is not supported as the control flow analyzer doesn't "see" the connection. For example:

type Data = { kind: 'str', payload: string } | { kind: 'num', payload: number };

function foo({ kind, payload }: Data) {
    if (kind === 'str') {
        payload.length;  // Error, payload not narrowed to string
    }
}

We may be able to support that pattern later, but likely not in this PR.

Answer 2

In TS 4.6 can use Wishlist: support for correlated union types #30581

type TypeA = {
  input: { foo: string }
  output: { foo: number }
}

type TypeB = {
  input: { bar: string }
  output: { bar: number }
}

type Types = {
  TypeA: TypeA
  TypeB: TypeB
}

type TypeMap = { 
  TypeA: TypeA, 
  TypeB: TypeB 
};

type TypeMapAsGeneric<K extends keyof TypeMap = keyof TypeMap> = { [P in K]: TypeMap[P] }[K];

function transformTypeA(typeA: TypeA['input']): TypeA['output'] {
  return { foo: parseInt(typeA.foo) }
}

function transformTypeB(typeB: TypeB['input']): TypeB['output'] {
  return { bar: parseInt(typeB.bar) }
}

const transformers: { [K in keyof TypeMap]: (data: TypeMapAsGeneric<K>['input']) => TypeMapAsGeneric<K>['output'] } = {
  TypeA: transformTypeA,
  TypeB: transformTypeB
};


const transform = <K extends keyof TypeMap>(typeName: K, inputValue: TypeMapAsGeneric<K>['input']): TypeMapAsGeneric<K>['output'] => {
  const transformer = transformers[typeName];
  return transformer(inputValue);
}


const transformedValue = transform('TypeA', { foo: '123' })
console.log(transformedValue)

Playground link