02 · 生成器与惰性求值

# 调用生成器函数，函数体一行都不执行
def gen():
    print("  [函数体开始跑了]")
    yield 1
    print("  [从暂停处恢复，继续往下]")
    yield 2
    print("  [函数体跑到底了]")

print("调用 gen() 之前")
g = gen()                       # ⭐ 这一行不执行函数体，只造一个生成器对象
print("调用 gen() 之后，拿到:", type(g).__name__)

print("第一次 next:", next(g))
print("第二次 next:", next(g))
try:
    next(g)
except StopIteration:
    print("第三次 next: 抛 StopIteration")

# yield 是暂停：局部变量和执行位置都被冻在生成器里
import inspect

def counter(start):
    i = start
    while i < start + 3:
        yield i
        i += 1                  # ⭐ 这一行要等【下一次 next】才执行

g = counter(10)
print("刚造出来      :", inspect.getgeneratorstate(g))
print("第 1 次 next ->", next(g), "| 冻住的 i =", g.gi_frame.f_locals["i"],
      "| 状态:", inspect.getgeneratorstate(g))
print("第 2 次 next ->", next(g), "| 冻住的 i =", g.gi_frame.f_locals["i"])
print("第 3 次 next ->", next(g), "| 冻住的 i =", g.gi_frame.f_locals["i"])
try:
    next(g)
except StopIteration:
    print("第 4 次 next -> StopIteration | 状态:", inspect.getgeneratorstate(g),
          "| gi_frame:", g.gi_frame)

# 生成器只能遍历一次，而且第二次【不报错】
g = (i * i for i in range(5))

print("第一次 sum :", sum(g))
print("第二次 sum :", sum(g))        # 不报错，安静地给 0
print("再 list    :", list(g), "| 长度", len(list(g)))

lst = [i * i for i in range(5)]
print("列表没这问题:", sum(lst), sum(lst))

# 生成器也没有 len，没有索引
try:
    len(x for x in range(5))
except TypeError as e:
    print("len(生成器) ->", type(e).__name__, ":", e)

# 生成器表达式 vs 列表推导：峰值内存差多少
import tracemalloc

N = 2_000_000

tracemalloc.start()
total = sum(i * i for i in range(N))            # ⭐ 生成器表达式：一次只活一个数
peak_gen = tracemalloc.get_traced_memory()[1]
tracemalloc.stop()

tracemalloc.start()
total2 = sum([i * i for i in range(N)])         # ⚠️ 列表推导：先把 200 万个数全造出来
peak_list = tracemalloc.get_traced_memory()[1]
tracemalloc.stop()

print(f"生成器表达式: 和={total}  峰值={peak_gen/1024/1024:.1f} MB")
print(f"列表推导    : 和={total2}  峰值={peak_list/1024/1024:.1f} MB")

# 惰性的代价：函数体里的参数校验也被推迟了
def read_rows(path):
    if not path.endswith(".csv"):
        raise ValueError("只支持 csv")     # ⚠️ 调用的那一刻【不会】执行
    yield 1

r = read_rows("data.txt")
print("调用完了，什么都没发生:", type(r).__name__)

try:
    next(r)
except ValueError as e:
    print("直到第一次 next 才炸:", type(e).__name__, ":", e)

# 想让校验立刻发生：外面一层普通函数，里面一层生成器
def read_rows_ok(path):
    if not path.endswith(".csv"):
        raise ValueError("只支持 csv")     # ⭐ 这是普通函数，调用即执行
    def _gen():
        yield 1
    return _gen()

try:
    read_rows_ok("data.txt")
except ValueError as e:
    print("调用时就炸:", type(e).__name__, ":", e)

# 迭代器协议：for 循环背后就这两个方法
class Countdown:
    def __init__(self, n):
        self.n = n
    def __iter__(self):
        return self             # ⚠️ 返回 self ＝ 这个对象【自己就是】迭代器，状态只有一份
    def __next__(self):
        if self.n <= 0:
            raise StopIteration
        self.n -= 1
        return self.n + 1

c = Countdown(3)
print("第一次 for:", [x for x in c])
print("第二次 for:", [x for x in c])        # ⚠️ 空的，状态已经耗尽

class CountdownOK:
    def __init__(self, n):
        self.n = n
    def __iter__(self):
        n = self.n
        while n > 0:            # ⭐ 每次 __iter__ 都造一个【新的】生成器
            yield n
            n -= 1

c2 = CountdownOK(3)
print("可重复遍历:", [x for x in c2], [x for x in c2])

# itertools 三件套：islice / chain / groupby
import itertools

def naturals():                  # 一个无穷流
    i = 0
    while True:
        yield i
        i += 1

print("islice 前 5 个:", list(itertools.islice(naturals(), 5)))
print("islice 第 5–8 :", list(itertools.islice(naturals(), 5, 9)))
print("chain 接起来  :", list(itertools.chain([1, 2], (3, 4), "ab")))

# ⚠️ groupby 只合并【相邻】的相同键 —— 不排序就得到一堆碎片
rows = [("a", 1), ("b", 2), ("a", 3)]
print("没排序:", [(k, [r[1] for r in grp])
                for k, grp in itertools.groupby(rows, key=lambda r: r[0])])
rows.sort(key=lambda r: r[0])
print("排序后:", [(k, [r[1] for r in grp])
                for k, grp in itertools.groupby(rows, key=lambda r: r[0])])

# 💀 groupby 的第二个坑：分组本身也是惰性的，往前走一步前一组就没了
print("先 list 再看:", [(k, list(grp)) for k, grp in list(itertools.groupby("aabbc"))])

# yield from：把子生成器的产出直接转出去
def leaves(node):
    if isinstance(node, list):
        for child in node:
            yield from leaves(child)      # ⭐ 一层 for 就够了
    else:
        yield node

tree = [1, [2, [3, 4], 5], [[6]]]
print("yield from :", list(leaves(tree)))

def leaves_manual(node):
    if isinstance(node, list):
        for child in node:
            for x in leaves_manual(child):   # 手写要两层 for
                yield x
    else:
        yield node

print("手写等价物 :", list(leaves_manual(tree)))

# ⚠️ yield from 和 for-yield 不等价的地方：它还转发【返回值】
def inner():
    yield 1
    yield 2
    return "inner 的返回值"

def outer():
    got = yield from inner()          # ⭐ StopIteration.value 落到 got 上
    print("  outer 收到:", got)

print("跑一遍     :", list(outer()))

# 读大文件：readlines() 把整个文件塞进内存，for line in f 不会
import os
import tracemalloc

PATH = "big_demo.txt"
with open(PATH, "w", encoding="utf-8") as f:
    for i in range(500_000):
        f.write(f"line {i} " + "x" * 40 + "\n")
print("文件大小:", round(os.path.getsize(PATH) / 1024 / 1024, 1), "MB")

tracemalloc.start()
n = 0
with open(PATH, encoding="utf-8") as f:
    for line in f:                      # ⭐ 文件对象本身就是迭代器，一次一行
        n += len(line)
peak_stream = tracemalloc.get_traced_memory()[1]
tracemalloc.stop()

tracemalloc.start()
with open(PATH, encoding="utf-8") as f:
    m = sum(len(line) for line in f.readlines())   # ⚠️ 先造一个 50 万元素的列表
peak_all = tracemalloc.get_traced_memory()[1]
tracemalloc.stop()

print(f"for line in f: 字符数={n} 峰值={peak_stream/1024/1024:.1f} MB")
print(f"readlines()  : 字符数={m} 峰值={peak_all/1024/1024:.1f} MB")
os.remove(PATH)