itertools.gropby使用例子

gropby, itertools

itertools-gropby-example

问题背景

本周主要写了各种类型的财务报表，利用python的csv模块进行输出。

在编写财务报表的过程中，出现了这样一种需求。

姓名 月份 电费 水费 燃气费
张三 1 100.00
张三 1 30.00
张三 1 40.00
张三 1 10.00
张三 2 70.00
李四 3 66.00
李四 3 20.00

类似于这样的报表，要将姓名和月份一致的信息，合并到一条中去。结果像这样：

姓名 月份 电费 水费 燃气费
张三 1 100.00 30.00 40.00
张三 2 70.00
李四 3 66.00 20.00

我们使用csv模块 的 DictWriter

表1对应的数据结构为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
[{
    'name': '张三',
    'month': 1,
    'electricity_fee': 100.00,
}, {
    'name': '张三',
    'month': 1,
    'water_fee': 30.00,
}, {
    'name': '张三',
    'month': 1,
    'gas_fee': 40.00,
}, {
    'name': '张三',
    'month': 1,
    'water_fee': 10.00,
}, {
    'name': '张三',
    'month': 2,
    'water_fee': 70.00,
}, {
    'name': '李四',
    'month': 3,
    'electricity_fee': 66.00,
}, {
    'name': '李四',
    'month': 3,
    'gas_fee': 20.00,
}]

表2：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
[{
    'name': '张三',
    'month': 1,
    'electricity_fee': 100.00,
    'water_fee': 40.00,
    'gas_fee': 40.00,
}, {
    'name': '张三',
    'month': 2,
    'water_fee': 70.00,
}, {
    'name': '李四',
    'month': 3,
    'electricity_fee': 66.00,
    'gas_fee': 20.00,
}]

如何用简介优雅的代码将表1的数据结构转化为表2？

1
2
3
4
5
6
7
8
9
10
11
12
13
from itertools import groupby
def get_key(data):
    return data['name'] + str(data['month'])
fees = ('electricity_fee', 'water_fee', 'gas_fee')
rows_out = []
for key, group in groupby(rows, key=get_key):
    row = group[0]
    for fee in fees:
        row[fee] = sum(item[fee] for item in group if fee in item)
        rows_out.append(row)