Inferring return type from the rest parameter shape

user_4964791 0
Viewed 78

So consider a set of transform functions that either take one argument and return a value or a function (of the same effect) which takes one or more arguments returning an array of values.

I don't want to implement a multi-function solution, e.g.

class A<T>  {
  transform(value: T): T { ... }
  transformMany(...values: T[]): T[] { ... }
}

An obvious solution is to fix the return as an array

Class A<T> {
  transform(...values: T[]): T[] { ... }
}
const a = new A<number>();
let one: number = a.transform(1)[0];
let many: number[] = a.transform(1,2,3);

but I find hardcoding the [0] unwise (and annoying).

Using infer yields a decent outcome

type TransformResultType<T> = T extends (infer R) ? R : T[];

class A<T>  {
  public transform<K extends T | T[] = T>(...values: T[]): TransformResultType<K> {
    const result: T[] = /* result of some transform */
    return (result.length ? result : result[0]) as TransformResultType<K>;
  }
}

const a: A<number> = new A<number>();
a.transform(1); // return type = number
a.transform<[]>(1, 2, 3); // return type = [] which is OK

However what I want is the return type to inferred by the shape of the call parameters rather than forcing the issue with <[]>.

What I want

a.transform(1); // return type = number
a.transform(1, 2, 3); // return type = number[]

Am I missing something or is this as good as it gets?

typescript
1 Answers

So you want a function that returns a single value with one argument, and an array of values with more than one argument?

You can do this by inferring the arguments as a tuple, and then testing its length in a conditional type.

Something like:

type TransformResultType<T extends unknown[]> =
  T['length'] extends 1 ? T[number] : T

// test it
type A = TransformResultType<[1]> // 1
type A = TransformResultType<[1, 2, 3]> // [1, 2, 3]

Note that the type is not number and number[] it is 1 and [1,2,3]. They are constant tuples, but they should be assignable to number and number[], so that shouldn't be a problem.

You now need to grab the whole array of arguments as the generic parameter, since you need the length. If you only infer the member type of the array, then you won't have that length.

class A<T>  {
  public transform<K extends T[]>(...values: K): TransformResultType<K> {
    const result: T[] = [] /* result of some transform */
    return (result.length ? result : result[0]) as TransformResultType<K>;
  }
}

Also note that the spread of arguments is always an array. So K extends T[] is the correct constraint for all arguments. T | T[] would allow for:

transform(1, 2, [4,5,6])

Which doesn't sound like what you want.

Playground

This approach does have some limitations, however. If you can't know the number of arguments at compile time, then you can't know the return type.

const data: number[] = [1]
const weird = a.transform(...data) // number[]

This would cause a problem, because the function is typed to return an array here, but the function will return a single value. The correct return value in this case is number | number[] because the type system doesn't actually know which it will be. Then the caller can sort it out.

To fix that we have to do a bit more work.

type Exact<A, B> =
  A extends B ? B extends A ? A : true : false

type IsLengthKnown<T extends unknown[]> =
  Exact<T['length'], number>

Exact is a handy type to test if two types are exactly the same. Exact<number, number> is true, but Exact<number, 1> is false.

This is important because the length of an array is number, but the length of a tuple is some specific number. So IsLengthKnown will tell us if T is a tuple with a fixed length or an array.

type TupleHasLength = IsLengthKnown<[1,2,3]> // true
type ArrayHasLength = IsLengthKnown<number[]> // false

Lastly, we can change the TranformResultType to this:

type TransformResultType<T extends unknown[]> =
  T['length'] extends 1 ? T[number] // one argument, use the member type of T
  : IsLengthKnown<T> extends true ? T // more than one argument use T
  : T | T[number] // uknown length, use a union of both

Which now does this:

const a: A<number> = new A<number>();

const single = a.transform(1); // 1
const multiple = a.transform(1, 2, 3); // [1,2,3]

const dataArray: number[] = [1]
const singleDataArray = a.transform(...dataArray) // number | number[]

const dataSingle = [1] as const
const singleDataTuple = a.transform(...dataSingle) // 1

I believe that will be correct in all cases.

Playground

All in all, always returning an array maybe be a better idea. As you can see it gets pretty hairy to handle all the edge cases. And as the consumer of this API, I think it's less cognitive load to know this function always accepts/returns an array, then to know there is a special case of one argument.

Squeegy_620760
Related