Generic zip in typescript, using variadc tuple-CodePudding

I have a background in python, and I enjoy playing with functional languages, and I'm learning typescript, so I'd like to have a zip() function that works somewhat like Python's. I see the typescript has gained both generators and variadic generics, so I think this should be possible. I have something that is very close to what I want, but it's not quite. I'm sure I'm either doing something obvious, or I'm about to learn the deep workings of the typescript type system.

What I want is a function that takes a variable number of arrays, and yields a tuple of elements from those arrays until the shortest array is exhausted. I wrote a generator that does that, but I'm having trouble getting the typing right.

  function* zip<T extends any[]>(...args: T) {
    for (let i = 0; i < Math.min(...args.map((e) => { return e.length }));   i) {
      yield args.map((e) => { return e[i]; }) as T;
    }
  }

If I remove the as T then the function works, but I lose all type information, if not then I get A return type [...T[]], but I want [...T].

For example, zip([1, 2, 3], ['a', 'b', 'c']), should have a return type of [number, string], but it has [number[], string[]]

CodePudding user response：

You probably want typings that look something like this:

function* zip<T extends any[][]>(...args: T) {
    for (let i = 0; i < Math.min(...args.map((e) => { return e.length }));   i) {
        yield args.map((e) => { return e[i]; }) as
            { [I in keyof T]: T[I] extends Array<infer E> ? E : never };
    }
}

Note that you want args to be an array of arrays, right? Each element of args should itself have a length property and have elements accessible via numeric indices. So T extends any[][] will do that for you (you might even want to loosen it up so that it accepts readonly arrays also, but I won't get into that.)

You don't want to return T for each element yielded by the generator, since that's an array of arrays. If you don't annotate the return then you get whatever the compiler thinks map() produces, which is necessarily going to be less precise than what you're looking for (see Mapping tuple-typed value to different tuple-typed value without casts for more information). You have to assert the type, but it's not T.

Instead, for each numeric index I in the indices of T, you want to take the array type T[I] at that index and return the element type of that array. We can use a mapped type to represent this transformation, because mapped types on tuples produce tuples. That's what { [I in keyof T]: T[I] extends Array<infer E> ? E : never } does.

Okay, let's test it:

for (const z of zip([1, 2, 3], ['a', 'b', 'c'])) {
    console.log(z[0].toFixed(2)   ", "   z[1].toUpperCase())
    // "1.00, A" "2.00, B" "3.00, C"
}

Looks good; the compiler sees that z is [number, string] and lets you treat each element of z accordingly.

Playground link to code