ndarray type annotation typescript

Question

I created a function that creates an array from a given shape with typescript. But how do I specify the return type of that function as it will vary depending on the passed shape argument? Is there a chance to express that with generics?

Examples:

createArray(shape: number[]): ???

>>> createArray([2, 3]) // return type: number[][]
[[0, 0, 0],
 [0, 0, 0]]

>>> createArray([2, 3, 2]) // return type: number[][][]
[[[0, 0],
  [0, 0], 
  [0, 0]],

 [[0, 0],
  [0, 0], 
  [0, 0]]]

As you can infer, the number of [] will always be the same as the length of the shape array. I assume you can't get the type annotation accurate because it is a dynamic type but is it at least possible to get a supertype, which is not any, that is variably deep nested?

Answer 1

Maybe this can help you:

type N_Arr<T extends any[]> = 
  T extends [any, any, ...any] 
    ? T extends [any, ...infer K] 
      ? N_Arr<K>[] 
      : number[] 
    : number[]

function createArray<T extends number[]>(shape: [...T]): N_Arr<T> {
  return [] as any
}

const t1 = createArray([0])           // t1: number[]
const t2 = createArray([0, 1])        // t2: number[][]
const t3 = createArray([0, 1, 2])     // t3: number[][][]
const t4 = createArray([0, 1, 2, 3])  // t4: number[][][][]

I created a type N_Arr which recursively adds [] to the type for each member in the tuple.

---

Explanation

This implementation relies on tuples. Tuples are similar to arrays but instead of being a list of some type of unknown length, each position in a tuple matters.

// array
type arr = string[]

// tuple
type tuple = [string, string]

As you can see above, the tuple is very restrictive and allows to specify type and position of each element. This can help us with the type of the createArray function because the length of the input array is essential for the return type.

The first step to the solution is to tell TypeScript that the input of the function has to be interpreted as a tuple type. We also want to store this tuple in a generic type T for further usage.

function createArray<T extends number[]>(shape: [...T]) /* ... */

This can be done by using the spread operator as shown above. If we don't do this and just declare shape as number[], the information about specific elements in the array is lost.

Now we need a generic type which takes the tuple T and constructs the nested array type. I called it N_Arr. N_Arr takes the tuple and does two things:

1. Check how many elements are left in the tuple.

2. Remove the first element of the tuple and call N_Arr recursively with the rest of the tuple.

For each any in the tuple another [] is appended to the type. The last recursive step evaluates to number[].

Playground