how to type define a zip function in typescript

Viewed 1407

Say I want to create a zip function:

function zip(arrays){
    // assume more than 1 array is given and all arrays 
    // share the same length
    const len = arrays[0].length;
    const toReturn = new Array(len);
    for (let i = 0; i < len; i++){
        toReturn[i] = arrays.map(array=>array[i]);
    }
    return toReturn;
}

console.log(zip([
    [1,2,3],
    [4,5,6],
    [7,8,9],
    [10,11,12],
]));
/*
Output:
(3) [Array(4), Array(4), Array(4)]
0: (4) [1, 4, 7, 10]
1: (4) [2, 5, 8, 11]
2: (4) [3, 6, 9, 12]
*/

In order to type define this function when all arrays hold the same type of elements:

    function zip<T>(arrays: T[][]): T[][]{/* codes omited here */}

However, when arrays are of different types of elements, I get confused about how to use generic type to finish the type definition.

    const zipedResult = zip([[1,2,3],[true,false,true],['a','b','c']]);
    // error raises: Type 'false' is not assignable to type 'number'.(2322)

what I want is

    [[1,2,3],[true,false,true],['a','b','c']]

could be automatically infered as (number|boolean|string)[][] without writing as (number|boolean|string)[][] or EVEN infered as [number[],boolean[],string[]] and result of zip infered as [number, boolean, string][]

How should I correctly type define zip to fullfill such features?

6 Answers

Here is an implementation that works for me:

export function zip<T extends unknown[][]>(
  ...args: T
): { [K in keyof T]: T[K] extends (infer V)[] ? V : never }[] {
  const minLength = Math.min(...args.map((arr) => arr.length));
  // @ts-expect-error This is too much for ts
  return range(minLength).map((i) => args.map((arr) => arr[i]));
}

for example zip(["x", "y", "z"], [true, false, true]) has inferred type [string, boolean][]

The solutions above are all valid and should work for normal use cases. We normally don't really know what elements are inside the array we would pass into the zip function. So the return type of

zip([1, 2, 3], ["a", "b", "c"])

being

[number, string][]

is fine for most cases where arbitrary arrays are passed into the function.

But if the elements are known at compile time, TypeScript allows us to have more accurate return types. We can change the defintion of zip, so that the return type of zip([1, 2, 3], ["a", "b", "c"]) is

[[1, "a"], [2, "b"], [3, "c"]]

Here is a solution that accomplishes this:

type ValidContent =
  | null
  | string
  | number
  | boolean
  | Array<JSON>
  | Date
  | undefined
  | {
    [prop: string]: ValidContent
  }

type ZipReturn<T extends any[][]> = T[0] extends infer A 
  ? { 
      [K in keyof A]: [...{ 
        [K2 in keyof T]: K extends keyof T[K2] ? T[K2][K] : undefined 
      }] 
    }
  : never

function zip<
  T extends [...{[K in keyof S]: S[K]}][], S extends (ValidContent)[]
>(...arrays: [...T]): ZipReturn<T> {
  const maxLength = Math.max(...arrays.map((x) => x.length));

  return range(maxLength).map((_, i) => {
    return range(arrays.length).map((_, k) => arrays[k][i]);
  }) as unknown as ZipReturn<T>;
}

Here are some test cases:

const a = zip([1, 2, 3])
//   ^? [[1], [2], [3]]

const b = zip([1, 2, undefined], ["a", "b", "c"])
//   ^? [[1, "a"], [2, "b"], [undefined, "c"]]

const c = zip([1, 2, 3], ["a", "b", "c"], [new Date(), new Date(), new Date()])
//   ^? [[1, "a", Date], [2, "b", Date], [3, "c", Date]]

const d = zip([1, 2, 3] as number[], ["a", "b", "c"] as string[])
//   ^? [number, string][]

As you can see in the last example, the function still works fine for arbitrary array types.

Playground

The only way I can see to do this is to define a different zip function for each size of array you'd like to handle (as we need to be able to say exactly what is in each part of the zip result):

const zip3 = <T, U, V>(arrays: [T[], U[], V[]]): [T, U, V][] => {
    const len = arrays[0].length;
    const toReturn: [T, U, V][] = new Array(len);
    for (let i = 0; i < len; i++){
        toReturn[i] = [arrays[0][i], arrays[1][i], arrays[2][i]];
    }
    return toReturn;
};

const result = zip3([
    [1,2,3],
    [true, false, true],
    [7,8,9],
]);

console.log(result);

Hopefully someone can come in and show a better way to do this, without having to redefine the function dependent on how many arrays you'd like to zip;

The most common solution would be:

declare type UnionTypes = number[] | string[] | boolean[];
function zip(arrays: UnionTypes[]): UnionTypes[]

I believe this is one of the few times when you see the complications of TS not being useful. I can only think of relaxing the type to just any and do manual type checking when necessary.

function zip(arrays: any[][]): any[][]{
    // assume more than 1 array is given and all arrays 
    // share the same length
    const len = arrays[0].length;
    const toReturn = new Array(len);
    for (let i = 0; i < len; i++){
        toReturn[i] = arrays.map(array => array[i]);
    }
    return toReturn;
}

const zipedResult = zip([[1,2,3],[true,false,true],['a','b','c']]);
console.log(zipedResult); // [[1, true, "a"], [2, false, "b"], [3, true, "c" ]]

Here's what I came up with:

  • fully typed
  • accepts arrays of different types
  • if the arrays do not have the same length, the result will have the length of the shorter array (which AFAIK is the most common implementation of zip)
function zip<A, B>(as: A[], bs: B[]): [A, B][]
{
    return as.length <= bs.length
         ? as.map((a, i) => [a, bs[i]])
         : bs.map((b, i) => [as[i], b])
}
Related