Python itertools.tee: แยก iterator ได้ แต่ต้องรู้ว่า buffer โตตรงไหน

Iterator อ่านได้รอบเดียว แต่บางงานต้องส่ง stream เดียวให้ consumer สองตัว itertools.tee สร้าง iterator อิสระหลายทางโดยอ่าน source ร่วมกันและเก็บ backlog ให้ทางที่ช้าครับ

from __future__ import annotations

import unittest
from itertools import tee


class TeeTests(unittest.TestCase):
    def test_source_is_pulled_once_and_slow_reader_catches_up(self) -> None:
        pulled: list[int] = []

        def source():
            for value in range(3):
                pulled.append(value)
                yield value

        fast, slow = tee(source(), 2)
        self.assertEqual(next(fast), 0)
        self.assertEqual(next(fast), 1)
        self.assertEqual(pulled, [0, 1])
        self.assertEqual(next(slow), 0)
        self.assertEqual(pulled, [0, 1])


if __name__ == "__main__":
    unittest.main()

ให้นึกถึงร้านที่ถ่ายสำเนาใบคิวให้แถวสองแถว ถ้าแถวแรกเดินเร็วกว่า ใบคิวที่แถวสองยังไม่อ่านต้องกองรออยู่ ความจำของ tee จึงขึ้นกับระยะห่างระหว่าง consumer ไม่ใช่แค่จำนวน iterator หากทางหนึ่งกิน stream จนจบ แต่อีกทางยังไม่เริ่ม buffer อาจโตใกล้ขนาดข้อมูลทั้งหมด

เอกสารยังระบุว่า tee ไม่ thread-safe อย่าให้ iterator ลูกถูกใช้พร้อมกันข้าม thread โดยไม่มีการประสาน หากรู้ว่าจะอ่านข้อมูลทั้งหมดหลายรอบและข้อมูลพอดี memory การสร้าง list(source) อาจตรงไปตรงมาและคาดการณ์ต้นทุนง่ายกว่า

Recap

• ใช้เมื่อ consumer เดินใกล้กันและ source สร้างซ้ำแพง
• เฝ้าระยะห่าง เพราะ backlog คือ memory
• ไม่ใช้ iterator ลูกพร้อมกันข้าม thread
• ใช้ list เมื่อจะเก็บทั้งหมดอยู่แล้ว