使用装饰器恢复发电机

时间:2014-10-30 14:30:30

标签: python python-3.x decorator yield

让我们有一个具有不时失败功能的类,但在某些操作之后它才能完美运行。

现实生活中的例子是提升_mysql_exceptions.OperationalError: (2006, 'MySQL server has gone away')的Mysql Query,但在客户端重新连接后它可以正常工作。

我试图为此编写装饰器:

def _auto_reconnect_wrapper(func):
    ''' Tries to reconnects dead connection
    '''

    def inner(self, *args, _retry=True, **kwargs):
        try:
            return func(self, *args, **kwargs)

        except Mysql.My.OperationalError as e:
            # No retry? Rethrow
            if not _retry:
                raise

            # Handle server connection errors only
            # http://dev.mysql.com/doc/refman/5.0/en/error-messages-client.html
            if (e.code < 2000) or (e.code > 2055):
                raise

            # Reconnect
            self.connection.reconnect()

        # Retry
        return inner(self, *args, _retry=False, **kwargs)
    return inner

class A(object):
    ...

    @_auto_reconnect_wrapper
    def get_data(self):
        sql = '...'
        return self.connection.fetch_rows(sql)

如果客户失去连接,它只是默默地重新连接,每个人都很高兴。

但是如果我想将get_data()转换为生成器(并使用yield语句)该怎么办:

    @_auto_reconnect_wrapper
    def get_data(self):
        sql = '...'
        cursor = self.connection.execute(sql)
        for row in cursor:
            yield row

        cursor.close()

嗯,之前的例子不会起作用,因为内部函数已经返回了生成器,并且在调用next()之后它将会中断。

据我所知,如果python在方法中看到yield,它只会立即产生控制(而不执行单个语句)并等待第一个next()

我设法通过替换:

使其成功
return func(self, *args, **kwargs)

使用:

for row in func(self, *args, **kwargs):
    yield row

但我很好奇是否有更优雅(更加pythonic)的方式来做到这一点。 有没有办法让python运行所有代码到第一个yield然后等待?

我知道只是调用return tuple(func(self, *args, **kwargs))的可能性,但我想避免一次性加载所有记录。

2 个答案:

答案 0 :(得分:6)

首先,我认为您目前使用的解决方案很好。当你装饰一个生成器时,装饰器至少需要像该生成器上的迭代器一样。通过使装饰器成为发电机来做到这一点也是完全可以的。正如x3al指出的那样,使用yield from func(...)代替for row in func(...): yield row是可能的优化。

如果你想避免实际上使装饰器成为一个生成器,你可以使用next来执行此操作,这将运行到第一个yield,并返回第一个产生的值。除了要由生成器生成的其余值之外,您还需要使装饰器以某种方式捕获并返回第一个值。您可以使用itertools.chain

执行此操作
def _auto_reconnect_wrapper(func):
    ''' Tries to reconnects dead connection
    '''

    def inner(self, *args, _retry=True, **kwargs):
        gen = func(self, *args, **kwargs)
        try:
            value = next(gen)
            return itertools.chain([value], gen)
        except StopIteration:
            return gen
        except Mysql.My.OperationalError as e:
            ...
            # Retry
            return inner(self, *args, _retry=False, **kwargs)
    return inner

您还可以使用inspect确定装饰器是否适用于生成器和非生成器函数,以确定您是否正在装饰生成器:

def _auto_reconnect_wrapper(func):
    ''' Tries to reconnects dead connection
    '''

    def inner(self, *args, _retry=True, **kwargs):
        try:
            gen = func(self, *args, **kwargs)
            if inspect.isgenerator(gen):
                value = next(gen)
                return itertools.chain([value], gen)
            else: # Normal function
                return gen
        except StopIteration:
            return gen
        except Mysql.My.OperationalError as e:
            ...
            # Retry
            return inner(self, *args, _retry=False, **kwargs)
    return inner

我赞成基于yield / yield from的解决方案,除非您需要除了生成器之外还要装饰常规函数。

答案 1 :(得分:3)

  

有没有办法让python运行所有代码直到第一次屈服然后再等待?

是的,它被称为next(your_generator)。只需拨打next()一次,代码就会在第一次yield之后完全等待。如果您不想丢失第一个值,可以在循环前放置另一个yield

如果你使用的是python 3.3+,你也可以替换

for row in func(self, *args, **kwargs):
    yield row

yield from func(self, *args, **kwargs)