Python的Generator和itertools

181 阅读4分钟
原文链接: zhuanlan.zhihu.com

关于generator和itertools的一些深度讨论。


要实现一个连接池。连接池经初始化建立,上层应用通过get_con这样的方法来获取连接。这里的连接池使用什么数据结构呢?

由于要对连接进行负载均衡,所以希望以O(1)的时间消耗均匀返回连接。这样看来,使用list最容易达成。但在归还(和释放)单个连接时,需要遍历查找连接池。

假设连接池中的元素有以下数据结构:

{
    'name': # the human readable name of the server
    'ip':  # the ip of the server
    'port': # the port of the server
    'con': # the connection itself, which we're going communicate on
    '...' #other fields like rtt
}

上层程序执行时,实际需要的是con对象,当发生异常时,需要通知连接池释放(或者销毁)这个对象,如果使用list数据结构来保存所有connection对象,就需要遍历list来查找。

考虑到也许连接断开的频率,这种遍历也许是可以接受的。但我们也有更好的方案。这就是使用dictionary来保存连接池数据结构,并用生成器来完成连接对象的负载均衡获取。这样,我们在获取和归还连接池对象时,都得到了O(1)的时间复杂度(存疑,可能取决于底层实现)。但无论如何,避免了归还时的O(n)时间复杂度。

def get_server(self):
    try:
        def foo():
            if not self.connection_pool:
                return None
            while True: #1
                for item in self.connection_pool.values(): #2
                    yield item #3
        self.get_server.__func__.iter = getattr(self.get_server.__func__, 'iter', foo()) #4
        return next(self.get_server.__func__.iter) #5
    except StopIteration:
        return None

这里截取的是某个连接类的方法。这个类有一个连接池connection_pool,类的其它方法在进行实际通信之前,通过get_server方法来获取一个连接。这里是返回整个server的数据结构,外面再通过decorator将实际连接传给调用者,这里略过。

1~3是常见的生成器语法。我们通过生成器,来保证每次调用get_server时,获取下一个连接,再通过外层的while True来循环往复。

重点是line 4~5。

如果我们每次调用get_server,都生成一个新的生成器序列的话,那么我们将永远只能获得连接池中第一个对象。这里需要一个其它语言里常见的函数静态变量,在python当中并不支持静态变量这种语法,我们需要另辟路径来实现。

在python中,函数也是一个对象,因此它可以持有自己的属性,并且由于函数的生命期是全局的,所以该属性的生命期也是全局的。这就满足了静态变量的第一个条件。另一方面,我们可以在访问函数时,通过一定的判断来给该静态变量初始化,从而使得该变量的初始化从代码封装的角度看,得以局部化,这就满足了静态变量的第二个条件。

下面是通过给函数增加自定义属性来模拟静态变量的一个例子:

def foo():
    print(foo.i)

foo.i = 0
foo()
===output===
0

但是,如果函数是类的成员方法,则不可以用上面的foo.i这样的方法来访问属性。而要用上面例子中的语法,即如果要给函数增加属性iter,需要通过self.foo.__func__.iter来访问。

另外值得一提的是,我们通常使用for ..in ..语法来访问生成器序列,这在get_server的使用场景下是不可用的。这里get_server()通过封装,使得我们可以以函数调用的方法,而不是for .. in 这样的语法来获取生成器序列中的一个值。

这里展开到另外一个问题。一般来说,对生成器序列,我们一旦获取其中的一个值,那么生成器序列的状态就发生了改变,如果我们要查看生成器中序列中的某个值,又不改变生成器状态,怎么办?即如何实现 peek功能。这一功能在跟踪时是很重要的。

这里需要用到itertools这个包。我们先从原先的生成器中取出(peek)一个元素,然后通过itertools.chain()来将它“还”回去。

from typing import Generator

def check_records(records: Generator):
    peek = next(records)
    print(peek)
    # return peek back
    return  itertools.chain([peek], records)