Python deep zip

Question

I am trying to write a function like zip. I am not good at explaining what I mean, so i will just show 'code' of what i'm trying to do.

a = [1,2,3,[4,5]]
b = a[:]
zip(a, b) == [(1,1), (2,2), (3,3), ([4,5],[4,5])]
myzip(a, b) == [(1,1), (2,2), (3,3), [(4,4), (5,5)]]

I am so stuck on this it's not even funny. I am trying to write it in a simple functional way with recursive lambdas, to make my code prettier. I want myzip like this because i want to use its output with another function I wrote which maps a function to a tree

def tree_map(func, tree):
    return map(lambda x: func(x) if not isinstance(x, list) else tree_map(func, x), 
               tree)

I have been trying to do something similar to this with zip, but I can't seem to wrap my head around it. Does anyone have any ideas on how i could write myzip?

Edit: Look at tree_map! isn't that pretty! i think so at least, but my mother tongue is Scheme :P and also, I want myzip to go as deep as it needs to. basically, I want myzip to retain the structure of the trees i pass it. Also, myzip will only handle trees that are the same shape.

Answer 1

I think the following should work:

import collections

def myzip(*args):
    if all(isinstance(arg, collections.Iterable) for arg in args):
        return [myzip(*vals) for vals in zip(*args)]
    return args

Result:

>>> a = [1,2,3,[4,[5,6]]]
>>> b = [1,2,3,[4,[5,6]]]
>>> myzip(a, b)
[(1, 1), (2, 2), (3, 3), [(4, 4), [(5, 5), (6, 6)]]]

Note that I use collections.Iterable instead of list in the type checking so that the behavior is more like zip() with tuples and other iterables.