使用装饰器恢复发电机 [英] Generator recovery using decorator

问题描述

让我们创建一个类，该类的功能会不时失败，但经过一些操作后，它会完美地发挥作用.

Let's have a class that has function that fails from time to time but after some actions it just works perfectly.

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

Real life example would be Mysql Query that raises _mysql_exceptions.OperationalError: (2006, 'MySQL server has gone away') but after client reconnection it works fine.

我试图为此编写装饰器:

I've tried to write decorator for this:

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)

如果客户端失去连接，它只会默默地重新连接，每个人都很高兴.

And if client loses connection it just silently reconnect and everybody is happy.

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

But what if I want to transform get_data() to generator (and use yield statement):

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

        cursor.close()

好吧，前面的示例不起作用，因为内部函数已经返回了generator，并且在调用第一个next()之后它将中断.

Well, previous example won't work because inner function already returned generator and it will break after calling first next().

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

As I understand it if python sees yield inside method it just yields control immediately (without executing one single statement) and waits for first next().

我设法通过替换使其起作用

I've managed to make it work by replacing:

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

使用:

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

但是我很好奇是否有更优雅(更pythonic)的方式来做到这一点. 是否有办法让python在所有yield和然后等待之前运行所有代码?

But I'm curious whether there is more elegant (more pythonic) way to do this. Is there a way to make python run all the code up to first yield and then wait?

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

I'm aware of possibility of just calling return tuple(func(self, *args, **kwargs)) but I want to avoid loading all records at once.

推荐答案

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

First, I think the solution you're currently using is fine. When you decorate a generator, the decorator is going to need to at least behave like an iterator over that generator. Doing that by making the decorator a generator, too, is perfectly ok. As x3al pointed out, using yield from func(...) instead of for row in func(...): yield row is a possible optimization.

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

If you want to avoid actually making the decorator a generator, too, you can do that by using next, which will run until the first yield, and return the first yielded value. You'll need to make the decorator somehow capture and return that first value, in addition to the rest of the values to be yielded by the generator. You could do that with 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 确定是否要装饰发电机:

You could also make the decorator work with both generator and non-generator functions, using inspect to determine if you're decorating a generator:

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的解决方案.

I would favor the yield/yield from-based solution, unless you have a requirement to decorate regular functions in addition to generators.

这篇关于使用装饰器恢复发电机的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！