使用装饰器恢复生成器

让我们有一个类，它的功能有时会失败，但经过一些操作后它就可以完美地工作。

现实生活中的例子是 Mysql 查询，它会引发_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()

好吧，前面的例子不起作用，因为内部函数已经返回生成器，并且它会在第一次调用后中断next().

据我了解，如果 python 看到yield在方法内部它只是立即产生控制（不执行任何一条语句）并等待第一个next().

我已经设法通过替换使其工作：

return func(self, *args, **kwargs)

With:

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

但我很好奇是否有更优雅（更Pythonic）的方法来做到这一点。有没有办法让python运行所有代码到第一个yield and then wait?

我知道可以打电话return tuple(func(self, *args, **kwargs))但我想避免一次加载所有记录。

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

如果您想避免实际上使装饰器也成为生成器，您可以通过使用来做到这一点next，它将运行到第一个yield，并返回第一个产生的值。除了生成器生成的其余值之外，您还需要使装饰器以某种方式捕获并返回第一个值。你可以这样做itertools.chain https://docs.python.org/3/library/itertools.html#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 https://docs.python.org/3/library/inspect.html确定您是否正在装饰发电机：

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基于 的解决方案，除非除了生成器之外您还需要装饰常规函数。