Python►中级进阶

常用标准库：itertools 与 functools

itertools 和 functools 是 Python 里两个”看起来抽象、实际非常实用”的标准库。前者提供各种迭代器组合器，后者提供函数式编程工具。两者搭配起来，你能用几行代码解决”分组”、”分块”、”缓存”、”偏函数”这些常见问题。这一篇讲清楚它们各自最有用的 15 个 API。

一、itertools：迭代器工具箱

无限迭代器

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from itertools import count, cycle, repeat

# 无限计数
for i in count(10, step=2):     # 10, 12, 14, ...
    if i > 20:
        break
    print(i)

# 无限循环
from itertools import islice
list(islice(cycle([1, 2, 3]), 8))       # [1, 2, 3, 1, 2, 3, 1, 2]

# 重复
list(repeat("x", 3))                    # ['x', 'x', 'x']

# repeat 常用来配合 map
list(map(pow, [1, 2, 3], repeat(2)))    # [1, 4, 9]  每个都平方

停止条件迭代器

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
from itertools import accumulate, chain, compress, dropwhile, takewhile, islice

# 累加（默认加法；可传 func）
list(accumulate([1, 2, 3, 4, 5]))        # [1, 3, 6, 10, 15]
list(accumulate([1, 2, 3, 4], func=max))  # [1, 2, 3, 4]  滑动最大
import operator
list(accumulate([1, 2, 3, 4], operator.mul))  # [1, 2, 6, 24]  阶乘

# 拼接多个可迭代对象
list(chain([1, 2], [3, 4], [5]))          # [1, 2, 3, 4, 5]
list(chain.from_iterable([[1, 2], [3, 4]]))  # 展平一层

# 掩码筛选（按位）
list(compress("ABCDEF", [1, 0, 1, 0, 1, 0]))  # ['A', 'C', 'E']

# 跳过直到条件为假
list(dropwhile(lambda x: x < 3, [1, 2, 3, 4, 1]))  # [3, 4, 1]

# 取到条件为假为止
list(takewhile(lambda x: x < 3, [1, 2, 3, 4, 1]))  # [1, 2]

# 切片
list(islice(range(100), 10, 20))          # [10..19]
list(islice(range(100), 5))                # [0..4]
list(islice(range(100), 0, None, 10))     # 每 10 个取一

组合迭代器

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from itertools import product, permutations, combinations, combinations_with_replacement

# 笛卡尔积
list(product([1, 2], "ab"))
# [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]

list(product(range(3), repeat=2))
# 所有 2 位数组合

# 排列
list(permutations("abc"))
# [('a','b','c'), ('a','c','b'), ('b','a','c'), ...]  共 3! = 6 个

list(permutations("abc", 2))       # 长度 2

# 组合（不考虑顺序）
list(combinations("abcd", 2))
# [('a','b'), ('a','c'), ('a','d'), ('b','c'), ('b','d'), ('c','d')]

# 允许重复
list(combinations_with_replacement("abc", 2))
# [('a','a'), ('a','b'), ('a','c'), ('b','b'), ('b','c'), ('c','c')]

用途：暴力枚举、算法题、生成测试用例。

分组

1
2
3
4
5
6
7
8
9
from itertools import groupby

data = [("A", 1), ("A", 2), ("B", 3), ("B", 4), ("A", 5)]

for key, group in groupby(data, key=lambda x: x[0]):
    print(key, list(group))
# A [('A', 1), ('A', 2)]
# B [('B', 3), ('B', 4)]
# A [('A', 5)]        ← 注意：只按"连续"分组

注意：groupby 只对相邻的相同 key 分组。要”全局分组”先 sorted：

1
2
3
data.sort(key=lambda x: x[0])
for key, group in groupby(data, key=lambda x: x[0]):
    print(key, list(group))

打包

1
2
3
4
5
6
7
8
9
10
11
12
13
14
from itertools import zip_longest, starmap, tee

# 以最长为准，短的补默认值
list(zip_longest([1, 2, 3], "ab", fillvalue="?"))
# [(1, 'a'), (2, 'b'), (3, '?')]

# 参数解包传给函数
list(starmap(pow, [(2, 3), (3, 2), (10, 2)]))
# [8, 9, 100]  相当于 [pow(2,3), pow(3,2), pow(10,2)]

# 分裂成 N 个独立迭代器
a, b = tee(range(5), 2)
list(a)     # [0, 1, 2, 3, 4]
list(b)     # [0, 1, 2, 3, 4]

注意：tee 底层会缓存已消费的元素——两个迭代器差距大时内存开销大，慎用。

batched（Python 3.12+）

1
2
3
4
from itertools import batched

list(batched("ABCDEFG", 3))
# [('A', 'B', 'C'), ('D', 'E', 'F'), ('G',)]

批量分块神器，3.12 之前要自己写（见 生成器 yield 深入）。

二、functools：函数式工具

partial / partialmethod：偏函数

前面 匿名函数 lambda 与高阶函数 讲过：

1
2
3
4
5
6
7
8
9
10
from functools import partial

def send(url, method="GET", data=None):
    print(url, method, data)

get = partial(send, method="GET")
post = partial(send, method="POST")

get("/api")           # /api GET None
post("/api", data={"a": 1})

reduce：累积计算

1
2
3
4
5
6
7
8
9
from functools import reduce

reduce(lambda a, b: a + b, [1, 2, 3, 4])          # 10
reduce(lambda a, b: a * b, [1, 2, 3, 4], 100)     # 2400，带初始值

# 用 operator 更快
import operator
reduce(operator.add, [1, 2, 3, 4])                 # 10
reduce(operator.mul, [1, 2, 3, 4])                 # 24

lru_cache / cache：函数缓存

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from functools import lru_cache, cache

@lru_cache(maxsize=128)
def fib(n):
    if n < 2:
        return n
    return fib(n-1) + fib(n-2)

fib(50)              # 秒出

# Python 3.9+ 无界版
@cache
def slow(x):
    ...

# 缓存统计
fib.cache_info()     # CacheInfo(hits=..., misses=..., maxsize=128, currsize=...)
fib.cache_clear()

要求：参数必须可哈希（list/dict/set 不行；元组/字符串/整数可以）。

wraps：保留元信息

前面 装饰器入门 讲过——写装饰器必备。

singledispatch：函数重载

Python 没有 Java 那种”同名不同参”重载，但用 singledispatch 可以按第一个参数的类型分派：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
from functools import singledispatch

@singledispatch
def render(obj):
    return f"未知类型：{obj}"

@render.register
def _(obj: int):
    return f"整数：{obj}"

@render.register
def _(obj: str):
    return f"字符串：{obj!r}"

@render.register(list)
def _(obj):
    return f"列表，长度 {len(obj)}"


print(render(3))            # 整数：3
print(render("hi"))         # 字符串：'hi'
print(render([1, 2, 3]))    # 列表，长度 3
print(render(None))         # 未知类型：None

在类方法里用 singledispatchmethod。

cached_property：类的懒加载属性

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from functools import cached_property

class Report:
    def __init__(self, data):
        self.data = data

    @cached_property
    def summary(self):
        print("计算中...")
        return sum(self.data)


r = Report([1, 2, 3])
r.summary        # 计算中... → 6
r.summary        # 直接返回 6，不再计算

第一次访问计算并存到实例 __dict__ 里，之后走属性查找。

total_ordering：补齐比较方法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from functools import total_ordering

@total_ordering
class Version:
    def __init__(self, major, minor):
        self.major, self.minor = major, minor

    def __eq__(self, other):
        return (self.major, self.minor) == (other.major, other.minor)

    def __lt__(self, other):
        return (self.major, self.minor) < (other.major, other.minor)


# 自动生成 <=, >, >=, !=
Version(1, 5) < Version(2, 0)    # True
Version(1, 5) >= Version(1, 5)   # True

三、operator 模块

不属于 functools，但配套使用非常多：

1
2
3
4
5
6
7
8
9
10
11
from operator import add, sub, mul, itemgetter, attrgetter, methodcaller

# 排序 key
sorted(pairs, key=itemgetter(1))
sorted(users, key=attrgetter("age"))
sorted(records, key=itemgetter("category", "price"))    # 多字段
sorted(names, key=methodcaller("lower"))

# 累加 / 累乘
reduce(add, [1, 2, 3])
reduce(mul, [1, 2, 3])

四、一些经典组合

分块（3.12 之前）

1
2
3
4
5
6
7
8
9
10
from itertools import islice

def chunked(iterable, size):
    it = iter(iterable)
    while chunk := tuple(islice(it, size)):
        yield chunk


list(chunked(range(10), 3))
# [(0, 1, 2), (3, 4, 5), (6, 7, 8), (9,)]

找连续差值

1
2
3
4
5
from itertools import pairwise    # 3.10+

for a, b in pairwise([1, 3, 6, 10]):
    print(b - a)
# 2 3 4

展平嵌套一层

1
2
3
4
from itertools import chain

flat = list(chain.from_iterable([[1, 2], [3, 4], [5]]))
# [1, 2, 3, 4, 5]

每个元素前面加 index

1
2
list(enumerate("abc"))                # 内置就行
list(zip(count(1), "abc"))            # 用 itertools 也行

五、一个完整例子：拼接多个 CSV

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
30
31
32
33
# combine.py
"""合并多个 CSV 文件，展示 itertools 的力量。"""

import csv
from pathlib import Path
from itertools import chain
from functools import reduce
from operator import iconcat


def read_rows(path: Path):
    """惰性地产出所有行。"""
    with path.open(encoding="utf-8", newline="") as f:
        reader = csv.reader(f)
        header = next(reader, None)
        for row in reader:
            yield row


def combine(paths: list[Path], out: Path, header: list[str]):
    """把多个 CSV 拼成一个。"""
    with out.open("w", encoding="utf-8", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(header)

        # chain 拼接所有文件的行，惰性写出
        for row in chain.from_iterable(read_rows(p) for p in paths):
            writer.writerow(row)


if __name__ == "__main__":
    files = sorted(Path("data").glob("part_*.csv"))
    combine(files, Path("all.csv"), ["id", "name", "score"])

这个例子的关键：chain.from_iterable(generator) 让我们一行行拼接，不用把全部数据读进内存。100GB 数据也能跑。

六、性能提示

  • itertools 都是惰性的——不消费不计算；
  • 大部分是 C 实现，比手写 Python 循环快；
  • functools.lru_cache 对递归 / 幂等的重计算是”免费加速”；
  • operator.itemgetter/attrgetter 比 lambda 快 20-50%。

七、常见陷阱

  1. groupby 只按连续分组：忘了 sort 会漏。
  2. tee 内存：两个迭代器差距大时缓存变大。
  3. lru_cache 参数不可哈希：TypeError。用 tuple 化，或改用别的缓存。
  4. reduce 空序列没初始值：TypeError。永远提供 initial 或保证非空。
  5. chain(a, b) vs chain.from_iterable([a, b])：一个直接传多个，一个传”一个装着可迭代对象的可迭代对象”。
  6. permutations 数量爆炸：10 个元素的全排列有 362 万种，别乱调。

八、小结与延伸阅读

  • itertools：无限迭代（count/cycle/repeat）、组合（product/permutations/combinations）、分组（groupby）、切片（islice）、拼接（chain）；
  • functools：partial / reduce / lru_cache / cached_property / singledispatch / total_ordering / wraps；
  • 配 operator：itemgetter / attrgetter 加速排序；
  • 处理大数据、组合数学、缓存、装饰器都能派上用场。

延伸阅读：

下一篇 日志模块 logging 实战 我们讲怎么写规范的日志。