python - 使用装饰器恢复发电机

让我们创建一个类，该类的功能有时会失败，但是经过一些操作后，它会完美地运行。

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

我试图为此编写装饰器:

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()

好吧，前面的示例将不起作用，因为内部函数已经返回了generator，并且在调用第一个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))的可能性，但我想避免一次加载所有记录。

最佳答案

首先，我认为您当前使用的解决方案很好。装饰生成器时，装饰器至少需要像该生成器上的迭代器一样工作。也可以通过使装饰器成为生成器来做到这一点。正如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的解决方案。

关于python - 使用装饰器恢复发电机，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/26655717/