Collate a list of dictionaries into a dictionary of lists

collate.py

__all__ = ['collate_dicts']

from typing import List, Dict, Any
from collections import defaultdict


def collate_dicts(ds: List[Dict[Any, Any]]) -> Dict[Any, List[Any]]:
    """
    Args:
        ds: A list of dictionarys all with the same keys.

    Returns:
        A dictionary of values collated across
    """
    if len(ds) == 0:
        return defaultdict(lambda: [])
    first_dict = ds[0]
    keys = first_dict.keys()
    collated = defaultdict(lambda: [])
    for d in ds:
        for k in keys:
            collated[k].append(d[k])
    return dict(collated)


def test_collate_dicts():
    assert collate_dicts([])['a'] == []
    assert collate_dicts([dict(a=1, b=2), dict(a=2, b=3)]) == dict(a=[1, 2], b=[2, 3])