TypeScript: infer return type of array based on input array elements

Question

I have a function which can have multiple parameters of type Node or Relationship or undefined. It has the following signature:

export function useFormat(...elements: Array<Node | Relationship | undefined>) {
    const formattedElements: Array<FormattedNode | FormattedRelationship | undefined> = [];
    // do formatting...
    return formattedElements;
 }

I call it as a React Hook and do some formatting based on the given element type. Then I return a new array with basically the same elements in a different format: Array<FormattedNode | FormattedRelationship | undefined>.

I would like to know if it is possible for TypeScript to infer the type of each element in the returned array based on the same element in the original array elements.

For example: the first element in the original array elements[0] is a Node, so first element of the returned array formattedElements[0] will be a FormattedNode.

captain-yossarian captain-yossarian · Accepted Answer · 2021-05-10T13:46:58

I believe you should use overloads here:

type Relationship = {
    type: 'Relationship'
}

type FormattedRelationship = {
    type: 'FormattedRelationship'
}

type CustomNode = {
    type: 'CustomNode '
}

type FormattedNode = {
    type: 'FormattedNode'
}


type Input = CustomNode | Relationship | undefined
type Output = FormattedNode | FormattedRelationship | undefined

type Return<T extends Input> =
    T extends CustomNode
    ? 0
    : T extends Relationship
    ? 1
    : T extends undefined
    ? 2
    : 3

function useFormat<
    T extends Input,
    U extends T[],
    R extends {
        0: FormattedNode[],
        1: FormattedRelationship[],
        2: undefined[],
        3: never
    }[Return<T>]>(...elements: [T, ...U]): R
function useFormat<
    T extends Input,
    U extends T[],
    >(...elements: [T, ...U]) {
    const formattedElements: Output[] = [];
    // do formatting...
    return formattedElements
}
const result = useFormat({ type: 'Relationship' }) // FormattedRelationship
const result2 = useFormat({ type: 'CustomNode ' }) // FormattedNode

Much readable format:


function useFormat<
    T extends CustomNode,
    U extends T[]>(...elements: [T, ...U]): FormattedNode[]
function useFormat<
    T extends Relationship,
    U extends T[]>(...elements: [T, ...U]): FormattedRelationship[]
function useFormat<
    T extends undefined,
    U extends T[]>(...elements: [T, ...U]): undefined[]
function useFormat<
    T extends Input,
    U extends T[],
    >(...elements: [T, ...U]) {
    const formattedElements: Output[] = [];
    // do formatting...
    return formattedElements
}
const result = useFormat({ type: 'Relationship' }) // FormattedRelationship
const result2 = useFormat({ type: 'CustomNode ' }) // FormattedNode

UPDATE

type Relationship = {
    type: 'Relationship'
}

type FormattedRelationship = {
    type: 'FormattedRelationship'
}

type CustomNode = {
    type: 'CustomNode '
}

type FormattedNode = {
    type: 'FormattedNode'
}


type Input = CustomNode | Relationship | undefined
type Output = FormattedNode | FormattedRelationship | undefined

type MapPredicate<T> =
    T extends Input
    ? T extends CustomNode
    ? FormattedNode
    : T extends Relationship
    ? FormattedRelationship
    : T extends undefined
    ? undefined
    : never : never

type Mapped<
    Arr extends Array<unknown>,
    Result extends Array<unknown> = []
    > = Arr extends []
    ? []
    : Arr extends [infer H]
    ? [...Result, MapPredicate<H>]
    : Arr extends [infer Head, ...infer Tail]
    ? Mapped<[...Tail], [...Result, MapPredicate<Head>]>
    : Readonly<Result>;

function useFormat<
    T extends Input,
    U extends T[],
    >(...elements: [...U]): Mapped<U>
function useFormat<
    T extends Input,
    U extends T[],
    >(...elements: [...U]): any {
    const formattedElements: Output[] = [];
    // do formatting...
    return formattedElements
}
const result = useFormat({ type: 'Relationship' }) // FormattedRelationship
const result2 = useFormat({ type: 'CustomNode ' }, { type: 'Relationship' }) // FormattedNode

Please keep in mind, these types are not 100% safe because I used any in order to make it compatible with overload.

Second, it is also better to define more than 1 overload. Since array types is mutable, some times it is hard to write type safe functions without using type assertions or any.

Now, you know all drawbacks.

Playground

If you want to know more about literal array mappings or tuples, you can refer to my blog

UPDATE - without rest

type Relationship = {
    type: 'Relationship'
}

type FormattedRelationship = {
    type: 'FormattedRelationship'
}

type CustomNode = {
    type: 'CustomNode '
}

type FormattedNode = {
    type: 'FormattedNode'
}


type Input = CustomNode | Relationship | undefined
type Output = FormattedNode | FormattedRelationship | undefined

type MapPredicate<T> =
    T extends Input
    ? T extends CustomNode
    ? FormattedNode
    : T extends Relationship
    ? FormattedRelationship
    : T extends undefined
    ? undefined
    : never : never

type Mapped<
    Arr extends ReadonlyArray<unknown>,
    Result extends Array<unknown> = []
    > = Arr extends []
    ? []
    : Arr extends [infer H]
    ? [...Result, MapPredicate<H>]
    : Arr extends [infer Head, ...infer Tail]
    ? Mapped<[...Tail], [...Result, MapPredicate<Head>]>
    : Readonly<Result>;

function useFormat<
    T extends Input,
    U extends ReadonlyArray<T>,
    >(elements: [...U]): Mapped<U>
function useFormat<
    T extends Input,
    U extends T[],
    >(elements: U): any {
    const formattedElements: Output[] = [];
    // do formatting...
    return formattedElements
}
const result = useFormat([{ type: 'Relationship' }]) // FormattedRelationship
const result2 = useFormat([{ type: 'CustomNode ' }, { type: 'Relationship' }]) // FormattedNode

Playground

This is the case when you need [...U] instead of U. It helps to infer every element in the array. Try to replace [...U] with U inside overload

TypeScript: infer return type of array based on input array elements

1 Answers