如何制作也可以用作上下文管理器的python方法？

斯图尔特·阿克森（Stuart Axon）

在我正在构建的图形API中，我想制作一个也可以用作contextmanager的方法：

该方法当前如下所示：

class Gfx:
    def __init__(self):
        self._fill = None
    
    def fill(self, color):
        self._fill = color

我可以创建一个contextmanager，将在使用后恢复状态，如下所示：

class Gfx:
    def __init__(self):
        self._fill = None

    @contextmanager
    def fill(self, color):
        """Allow user to set fill color, then restore it"""
        old_fill = self._fill
        self._fill = color 
        yield
        self._fill = old_fill

但是，如何根据它的调用方式使它同时工作呢？

>>> gfx = Gfx()
>>> gfx.fill("red") # use as ordinary method
>>> print(self._fill)
"red"

>>> with gfx.fill("blue") # use as context manager
...     print(gfx._fill)
"blue"
>>> print(gfx._fill)
"red"

暗影游侠

您需要混合使用各种方法。有一个称为的函数，返回一个上下文管理器，当不使用时，将忽略上下文管理器with。类似于以下内容：

class Gfx:
    def __init__(self):
        self._fill = None

    @contextmanager
    def _fillctx(self, old_fill):
        try:
            yield
        finally:
            self._fill = old_fill

    def fill(self, color):
        """Allow user to set fill color, then restore it"""
        old_fill = self._fill
        self._fill = color
        return self._fillctx(old_fill)

如果不使用调用with，self._fill则会设置并返回一个从未使用过的上下文管理器。当使用调用时with，它将在调用该上下文管理器的self._fill时重置__exit__。