迭代器与生成器

# for 循环写法（你平时写的）
for item in [1, 2, 3]:
    print(item)

# 等价的手动迭代（for 循环的底层实现）
iterable = [1, 2, 3]
iterator = iter(iterable)          # 调用 __iter__()
while True:
    try:
        item = next(iterator)      # 调用 __next__()
        print(item)
    except StopIteration:
        break                       # 迭代结束

class Countdown:
    """一个倒计时迭代器，每次 next() 返回递减的数字"""

    def __init__(self, start: int):
        self.current = start

    def __iter__(self):
        return self  # 迭代器返回自身

    def __next__(self):
        if self.current <= 0:
            raise StopIteration
        self.current -= 1
        return self.current + 1  # 返回递减前的值


# 使用示例
cd = Countdown(3)
print(list(cd))  # [3, 2, 1]

# 可以放在 for 循环中
for num in Countdown(3):
    print(num)
# 输出: 3 2 1

class CountdownIterable:
    """可迭代对象：每次 iter() 都返回一个新的迭代器"""

    def __init__(self, start: int):
        self.start = start

    def __iter__(self):
        return CountdownIterator(self.start)


class CountdownIterator:
    """迭代器：持有遍历状态"""

    def __init__(self, start: int):
        self.current = start

    def __iter__(self):
        return self

    def __next__(self):
        if self.current <= 0:
            raise StopIteration
        self.current -= 1
        return self.current + 1


# 现在可以重复遍历了
cd = CountdownIterable(3)
print(list(cd))  # [3, 2, 1]
print(list(cd))  # [3, 2, 1] —— 每次 iter() 创建新的迭代器

def countdown_gen(start: int):
    """生成器函数版本的倒计时"""
    while start > 0:
        yield start
        start -= 1


# 调用生成器函数，返回生成器对象
gen = countdown_gen(3)
print(type(gen))  # <class 'generator'>
print(list(gen))  # [3, 2, 1]

def sub_generator():
    yield 1
    yield 2
    yield 3


def main_generator():
    yield "start"
    yield from sub_generator()  # 委托给子生成器
    yield "end"


print(list(main_generator()))
# ['start', 1, 2, 3, 'end']

def accumulating_sub():
    """子生成器：接收 send 的值并累加"""
    total = 0
    while True:
        received = yield total
        if received is None:
            break
        total += received
    return total  # StopIteration 的 value


def accumulator():
    """主生成器：通过 yield from 获取子生成器的返回值"""
    result = yield from accumulating_sub()
    yield f"最终结果: {result}"


acc = accumulator()
print(next(acc))     # 0（启动子生成器，返回初始 total）
print(acc.send(10))  # 10（发送 10，累加后返回）
print(acc.send(20))  # 30（发送 20，累加后返回）
print(acc.send(None))  # "最终结果: 30"（结束子生成器，获取返回值）

# 列表推导式：立即计算，返回 list
squares_list = [x ** 2 for x in range(10)]
print(type(squares_list))  # <class 'list'>
print(squares_list)        # [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

# 生成器表达式：惰性计算，返回 generator
squares_gen = (x ** 2 for x in range(10))
print(type(squares_gen))   # <class 'generator'>
print(list(squares_gen))   # [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

import sys

# 列表推导式：一次性在内存中创建 1000 万个元素
# 请谨慎运行，这会占用大量内存
# big_list = [x ** 2 for x in range(10_000_000)]

# 生成器表达式：仅占用生成器对象本身的内存（约 200 字节）
big_gen = (x ** 2 for x in range(10_000_000))
print(f"生成器对象大小: {sys.getsizeof(big_gen)} bytes")
# 输出约: 生成器对象大小: 200 bytes

# 而等价的列表将占用约 80 MB
# print(f"列表大小: {sys.getsizeof(big_list)} bytes")  # 约 80,000,000 bytes

# 标准写法
total = sum((x ** 2 for x in range(100)))

# 简写（省略生成器表达式的括号）
total = sum(x ** 2 for x in range(100))

# 等价于（但更省内存）
total = sum([x ** 2 for x in range(100)])

def echo_generator():
    """回显生成器：接收值并返回"""
    print("生成器已启动")
    while True:
        received = yield  # yield 在等号右边，接收 send 的值
        print(f"收到: {received}")


echo = echo_generator()
next(echo)  # 必须先启动生成器（执行到第一个 yield）
# 输出: 生成器已启动

echo.send("Hello")  # 输出: 收到: Hello
echo.send("World")  # 输出: 收到: World
echo.close()

def resilient_generator():
    """能处理异常的生成器"""
    for i in range(5):
        try:
            yield i
        except ValueError as e:
            print(f"捕获到 ValueError: {e}")
            yield "error_handled"


gen = resilient_generator()
print(next(gen))  # 0
print(next(gen))  # 1
print(gen.throw(ValueError, "测试异常"))  # 捕获到 ValueError: 测试异常
                                         # error_handled
print(next(gen))  # 2（继续正常迭代）

def cleanup_generator():
    """带清理逻辑的生成器"""
    try:
        yield 1
        yield 2
    finally:
        print("清理资源：关闭数据库连接、释放文件句柄等")


gen = cleanup_generator()
print(next(gen))  # 1
gen.close()       # 在 yield 处抛出 GeneratorExit，触发 finally
# 输出: 清理资源：关闭数据库连接、释放文件句柄等

from itertools import count, cycle, repeat

# count(start, step): 从 start 开始，无限递增
counter = count(10, 2)
print([next(counter) for _ in range(5)])  # [10, 12, 14, 16, 18]

# cycle(iterable): 无限循环一个可迭代对象
cycler = cycle("ABC")
print([next(cycler) for _ in range(6)])   # ['A', 'B', 'C', 'A', 'B', 'C']

# repeat(obj, times): 重复一个对象（times 省略则无限重复）
repeater = repeat("hello", 3)
print(list(repeater))  # ['hello', 'hello', 'hello']

from itertools import accumulate, chain, pairwise, groupby

# accumulate: 累加（可自定义二元函数）
print(list(accumulate([1, 2, 3, 4])))  # [1, 3, 6, 10]

# chain: 串联多个可迭代对象
print(list(chain("AB", "CD", [1, 2])))  # ['A', 'B', 'C', 'D', 1, 2]

# pairwise: 相邻元素配对（Python 3.10+）
print(list(pairwise("ABCD")))  # [('A', 'B'), ('B', 'C'), ('C', 'D')]

# groupby: 按 key 分组（注意：需要先排序）
data = [("a", 1), ("a", 2), ("b", 3), ("b", 4)]
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)]

from itertools import product, permutations, combinations, combinations_with_replacement

# product: 笛卡尔积
print(list(product("AB", [1, 2])))
# [('A', 1), ('A', 2), ('B', 1), ('B', 2)]

# permutations: 排列（顺序重要）
print(list(permutations("ABC", 2)))
# [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]

# combinations: 组合（顺序不重要，无重复）
print(list(combinations("ABC", 2)))
# [('A', 'B'), ('A', 'C'), ('B', 'C')]

# combinations_with_replacement: 有放回组合
print(list(combinations_with_replacement("AB", 2)))
# [('A', 'A'), ('A', 'B'), ('B', 'B')]

def data_generator():
    for i in range(3):
        yield i


gen = data_generator()
print(list(gen))  # [0, 1, 2]
print(list(gen))  # [] —— 陷阱！

# 正确做法：每次需要遍历时重新创建
print(list(data_generator()))  # [0, 1, 2]
print(list(data_generator()))  # [0, 1, 2]

# 错误做法
data = [1, 2, 3, 4, 5]
for item in data:
    if item % 2 == 0:
        data.remove(item)  # 危险！跳过了一些元素
print(data)  # [1, 3, 5] —— 看似正确，但内部跳过了检查

# 正确做法
data = [1, 2, 3, 4, 5]
data = [item for item in data if item % 2 != 0]
print(data)  # [1, 3, 5]

def read_large_file(file_path: str):
    """逐行读取大文件，每次只在内存中保留一行"""
    with open(file_path, "r", encoding="utf-8") as f:
        for line in f:  # 文件对象本身就是可迭代的
            yield line.strip()


def filter_log_errors(file_path: str):
    """从日志文件中筛选包含 ERROR 的行"""
    for line in read_large_file(file_path):
        if "ERROR" in line:
            yield line


# 使用管道式处理
for error_line in filter_log_errors("/var/log/app.log"):
    print(error_line)

def read_data(source):
    """步骤1：读取数据"""
    for item in source:
        yield item


def filter_valid(records):
    """步骤2：过滤无效记录"""
    for record in records:
        if record.get("status") == "valid":
            yield record


def transform(records):
    """步骤3：转换数据格式"""
    for record in records:
        yield {
            "id": record["id"],
            "name": record["name"].upper(),
            "score": record["score"] * 1.1,
        }


def save_results(records, threshold: float):
    """步骤4：保存高分记录"""
    for record in records:
        if record["score"] > threshold:
            yield record


# 管道串联：每一步都是惰性的
raw = [
    {"id": 1, "name": "alice", "score": 85, "status": "valid"},
    {"id": 2, "name": "bob", "score": 60, "status": "invalid"},
    {"id": 3, "name": "charlie", "score": 92, "status": "valid"},
]

pipeline = save_results(
    transform(
        filter_valid(
            read_data(raw)
        )
    ),
    threshold=90.0,
)

# 直到这里才开始实际计算
for result in pipeline:
    print(result)
# {'id': 1, 'name': 'ALICE', 'score': 93.5}
# {'id': 3, 'name': 'CHARLIE', 'score': 101.2}

def fibonacci():
    """生成无限斐波那契数列"""
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b


# 取前 10 个斐波那契数
from itertools import islice

fib = fibonacci()
first_10 = list(islice(fib, 10))
print(first_10)  # [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

def paginated_fetch(api_client, page_size: int = 100):
    """分页获取 API 数据，对调用者透明"""
    page = 1
    while True:
        data = api_client.fetch(page=page, size=page_size)
        if not data:
            break
        yield from data  # 逐个产出记录
        page += 1


# 调用者不需要关心分页逻辑
for record in paginated_fetch(client, page_size=50):
    process(record)

# 不推荐：嵌套循环
def flatten_bad(nested):
    for sublist in nested:
        for item in sublist:
            yield item

# 推荐：使用 yield from
def flatten_good(nested):
    for sublist in nested:
        yield from sublist

# 读取文件直到遇到结束标记
with open("data.txt", "r") as f:
    for line in iter(lambda: f.readline().strip(), "END"):
        print(line)

from contextlib import closing


def db_query_generator(connection, query):
    cursor = connection.cursor()
    cursor.execute(query)
    try:
        for row in cursor:
            yield row
    finally:
        cursor.close()


# 使用 closing 确保即使提前退出也会清理
with closing(db_query_generator(conn, "SELECT * FROM users")) as results:
    for row in results:
        if row["id"] > 1000:
            break  # 提前退出，closing 确保 cursor.close() 被调用

# 不推荐：手写滑动窗口
def sliding_window_bad(seq, n):
    for i in range(len(seq) - n + 1):
        yield seq[i:i + n]

# 推荐：使用 itertools.islice 和 tee
from itertools import islice, tee


def sliding_window_good(iterable, n):
    iterators = tee(iterable, n)
    for i, it in enumerate(iterators):
        next(islice(it, i, i), None)
    return zip(*iterators)