摄影：产品经理一起吃：kingname & 产品经理
在上一篇文章中，我们说到了，itertools.tee不是线程安全的，并给出了一个例子，如下图所示：
在两个线程里面同时运行分裂出来的生成器对象，就会导致报错。
现在，你想看看itertools.tee的源代码，但是你会发现，在 PyCharm 里面，它的源代码如下图所示：
这是因为，在 CPython 中，itertools.tee底层是通过 C 语言实现的，所以你不能在 PyCharm 中看到它的源代码。但是你可以通过阅读 Python 的源代码中的 Modules/itertoolsmodule.c 文件[1]，找到它的实现算法。
导致问题的核心部分在如下图所示的两段代码中：
大家看不懂也没有关系，根据我上一篇文章中使用 Python 实现的简化版本就足够帮助理解了。
我们使用简化版本来解释其中线程不安全的地方：
def generator(): for i in range(3): yieldf'我是你第{i}个爷爷' def split(g): value_list_1 = [] value_list_2 = [] def wrap(queue): whileTrue: ifnot queue: try: value = next(g) except StopIteration: return value_list_1.append(value) value_list_2.append(value) yield queue.pop(0) g_1 = wrap(value_list_1) g_2 = wrap(value_list_2) return g_1, g_2 g = generator() g_1, g_2 = split(g) for value in g_1: print(value) for value in g_2: print(value)
当两个线程同时运行到if not queue时，发现当前各自的队列都是空的，于是进入value = next(g)获取下一个值。其中，线程 A 先进入那么几毫秒。然后线程 B 进入value = next(g)。但由于此时线程 A 中的next(g)正在运行，尚未结束，线程 B 又跑来运行，于是就导致了报错的发生。Python 中，生成器不是线程安全的。
那么如何让itertools.tee分裂出来的多个生成器可以在多线程中运行呢？其关键因素就是让value = next(g)这一行一次只能让一个线程运行。所以我们可以通过加锁来实现。
import itertools from threading import Lock class KingnameTee: def __init__(self, tee_obj, lock): self.tee_obj = tee_obj self.lock = lock def __iter__(self): return self def __next__(self): with self.lock: return next(self.tee_obj) def __copy__(self): return KingnameTee(self.tee_obj.__copy__(), self.lock) def safe_tee(iterable, n=2): """tuple of n independent thread-safe iterators""" lock = Lock() return tuple(KingnameTee(tee_obj, lock) for tee_obj in itertools.tee(iterable, n))
我们来看看运行效果：
多线程完美运行。
