I have a more general implementation that can accept any number of containers of any type as parameters.
from collections.abc import Iterable
import types
def dict_value_map(fun, *dicts):
keys = dicts[0].keys()
for d in dicts[1:]:
assert d.keys() == keys
return {k:fun(*(d[k] for d in dicts)) for k in keys}
def collection_map(fun, *collections):
assert len(collections) > 0
if isinstance(collections[0], dict):
return dict_value_map(fun, *collections)
if isinstance(collections[0], (tuple, list, set)):
return type(collections[0])(map(fun, *collections))
else:
return map(fun, *collections)
iscollection = lambda v:(isinstance(v,Iterable)and(not isinstance(v,str)))
def apply(fun, *collections, at=lambda collections: not iscollection(collections[0])):
'''
like standard map, but can apply the fun to inner elements.
at: a int, a function or sometype.
at = 0 means fun(*collections)
at = somefunction. fun will applied to the elements when somefunction(elements) is True
at = sometype. fun will applied to the elements when elements are sometype.
'''
if isinstance(at, int):
assert at >= 0
if at == 0:
return fun(*collections)
else:
return collection_map(lambda *cs:apply(fun, *cs, at=at-1), *collections)
if isinstance(at, types.FunctionType):
if at(collections):
return fun(*collections)
else:
return collection_map(lambda *cs:apply(fun, *cs, at=at), *collections)
else:
return apply(fun, *collections, at=lambda eles:isinstance(eles[0], at))
examples:
> apply(lambda x:2*x, [(1,2),(3,4)])
[(2, 4), (6, 8)]
> apply(lambda a,b: a+b, ([1,2],[3,4]), ([5,6],[7,8]))
([6, 8], [10, 12])
> apply(lambda a,b: a+b, ([1,2],[3,4]), ([5,6],[7,8]), at=1)
([1, 2, 5, 6], [3, 4, 7, 8])
> apply(lambda a,b: a+b, ([1,2],[3,4]), ([5,6],[7,8]), at=0)
([1, 2], [3, 4], [5, 6], [7, 8])
> apply(lambda a,b:a+b, {'m':[(1,2),[3,{4}]], 'n':5}, {'m':[(6,7),[8,{9}]],'n':10})
{'m': [(7, 9), [11, {13}]], 'n': 15}
> apply(str.upper, [('a','b'),('c','d')], at=str)
[('A', 'B'), ('C', 'D')]
and
> apply(lambda v:v+7, {'a': 1, 'b': {'c': 6, 'd': 7, 'g': {'h': 3, 'i': 9}}, 'e': {'f': 3}})
{'a': 8, 'b': {'c': 13, 'd': 14, 'g': {'h': 10, 'i': 16}}, 'e': {'f': 10}}