Skip to content

Instantly share code, notes, and snippets.

@zastari
Last active May 4, 2018 01:22
Show Gist options
  • Select an option

  • Save zastari/fb79b6b4c9fa617b6f00e8e31b8d2a49 to your computer and use it in GitHub Desktop.

Select an option

Save zastari/fb79b6b4c9fa617b6f00e8e31b8d2a49 to your computer and use it in GitHub Desktop.
groupby + rollup for namedtuples
def multigroup(entries, items=None, groups=None, rollup=True):
"""Object grouping with multiple groups that supports optional rollup
:arg entries: A list of objects that supports `getattr` for each field in `items`
:arg items: Attribute fields to retrieve from entry namedtuples
:arg groups: Attribute fields to group by
:arg rollup: Whether to put common fields in a group in "_rollup"
:type entries: list
:type items: list or None
:type groups: list or None
:type rollup: bool
:return: Group attributes nested in `groups` order.
Special fields::
_rollup: K:V pairs that were all identical inside of a given group if `rollup` = True
_entries: Entry dictionaries where keys are fields from `items` that weren't in `groups`
and weren't rolled up.
:rtype: dict
"""
group, groups = (groups[0], groups[1:]) if groups else (None, list())
items = [i for i in items if i != group] if items else list()
result = OrderedDict()
if rollup:
_rollup = {
i: getattr(entries[0], i) for i in items
if i not in groups and all(getattr(e, i) == getattr(entries[0], i) for e in entries)
}
if _rollup:
result['_rollup'] = _rollup
items = [i for i in items if i not in _rollup]
if not group:
result['_entries'] = [dict(s) for s in set(tuple((i, getattr(e, i)) for i in items) for e in entries)]
return result
keyfunc = lambda x: getattr(x, group)
for key, entry_group in groupby(sorted(entries, key=keyfunc), keyfunc):
result[key] = multigroup(list(entry_group), items, groups=groups, rollup=rollup)
return result
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment