在单个装饰器中处理不同类型的 return 值

Handle different kinds of return values in a single decorator

我有一个 Python class 有几个 class 方法,其中一些 return 一个 value 和其他 return (value1, value2)。他们应该return None/(None, None) 对于这些方法开头的一些错误,例如:

class Demo:
    def foo(*args, **kwargs) -> int:
        if not self._db.is_connected():
            return None
        # do something
        return x

    def bar(*args, **kwargs) -> Tuple[int, int]:
        if not self._db.is_connected():
            return (None, None)
        # do something
        return (x, y)

我尝试使用装饰器进行此检查。我知道这行得通:

class Demo:
    def _is_connected_decorator(func):
        def magic(self, *args, **kwargs):
            return func(self, *args, is_connected=self._db.is_connected(), **kwargs)
        return magic

    @_is_connected_decorator
    def foo(*args, is_connected=False, **kwargs) -> int:
        if not is_connected:
            return None
        # do something
        return x

    @_is_connected_decorator
    def bar(*args, is_connected=False, **kwargs) -> Tuple[int, int]:
        if not is_connected:
            return (None, None)
        # do something
        return (x, y)

但是,我想完全将支票放入单个装饰器中,而不修改 foo() 和 bar() 的 args 或 return 值,如下所示:

class Demo:
    def _is_connected_decorator(func):
        def magic(self, *args, **kwargs):
            if not self._db.is_connected():
                return None   # error return
            return func(self, *args, **kwargs)
        return magic

    @_is_connected_decorator
    def foo(*args, **kwargs) -> int:
        # do something
        return x

    @_is_connected_decorator
    def bar(*args, **kwargs) -> Tuple[int, int]:
        # do something
        return (x, y)

然而,装饰器并不知道它应该returnNone(None, None)。有没有好的方法来处理这两种 return 值?

可以指定默认值吗?

class Demo:
    def _is_connected_decorator(default):
        def decorator(func):
            def magic(self, *args, **kwargs):
                if not self._db.is_connected():
                    return default
                return func(self, *args, **kwargs)
            return magic
        return decorator

    @_is_connected_decorator(None)
    def foo(*args, **kwargs) -> int:
        # do something
        return x

    @_is_connected_decorator((None, None))
    def bar(*args, **kwargs) -> Tuple[int, int]:
        # do something
        return (x, y)

Demo at repl.it 使用此测试:

x, y = 1, 2
d = Demo()
d._db = type('', (), {})
d._db.is_connected = lambda: connected
for connected in False, True:
    print(d.foo())
    print(d.bar())

输出:

None
(None, None)
1
(1, 2)