Python：标准函数和上下文管理器？

Question

在 Python 中，有许多既可以作为标准函数又可以作为上下文管理器的函数。例如，open()可以这样调用：

my_file=open(filename,'w')

或者

with open(filename,'w') as my_file:

两者都会给你一个my_file对象，可用于进行所需的操作。通常情况下，后者更可取，但有时候也需要使用前者。

我已经能够学会如何编写上下文管理器了，可以通过创建一个带有__enter__和__exit__函数的类或者使用@contextlib.contextmanager装饰器在函数上并使用 yield 而不是 return 来实现。但是，当我这样做时，我不能再直接使用该函数——例如，使用装饰器，我得到了一个_GeneratorContextManager对象而不是想要的结果。当然，如果我将其作为类来制作，我将只获得生成器类的实例，我认为本质上它们是相同的东西。

那么我该如何设计一个既可以作为函数（返回对象），又可以作为上下文管理器（返回_GeneratorContextManager或类似对象）的函数（或类）呢？

编辑：

例如，假设我有一个像下面这样的函数（这是极度简化的）：

def my_func(arg_1,arg_2):
    result=arg_1+arg_2
    return my_class(result)

因此，该函数接受一些参数，对它们进行操作，并使用这些操作的结果来初始化一个类，然后返回它。最终结果是我拥有了一个my_class实例，就像我如果调用open一样拥有一个file对象。如果我想将这个函数作为上下文管理器使用，我可以这样修改它：

@contextlib.contextmanager
def my_func(arg_1,arg_2):
    result=arg_1+arg_2 # This is roughly equivalent to the __enter__ function
    yield my_class(result)
    <do some other stuff here> # This is roughly equivalent to the __exit__function

当作为上下文管理器调用时，它可以正常工作，但是当直接调用时，我不再获取my_class的实例。也许我只是做错了什么？

编辑2：

请注意，我对my_class有完全控制权，包括添加函数的能力。从下面的接受答案中，我能够推断出我的困难源于一个基本的误解：我认为我调用的任何内容（例如上面的my_func）都需要具有 __exit__和__enter__函数。这是不正确的。事实上，只有作为上下文管理器使用的返回结果（如上面的my_class）需要这些函数。

- ibrewster

---

open 是一个函数，它返回一个类的实例。无论您使用 myfile = open(filename) 还是 with open(filename) as myfile，它仍然是相同类的一个实例。没有任何改变。 - zondo

@zondo 是的，正如我试图编写的函数一样。但是，当我将该函数包装在 @contextlib.contextmanager 装饰器中，并像调用标准函数一样调用它时，我得到的类不是我从函数中“yield”的类。只有当我将其作为上下文管理器调用时，才会得到那个类。我将添加一个简单的示例。 - ibrewster

在你的类中定义一个__call__方法。 - apex-meme-lord

@apex-meme-lord：那么我应该编写一个具有__call__方法的上下文管理器类，该方法返回我想要的实例化类吗？如果我没有显式地使用__init__方法，这是否意味着当我执行my_instance=cm_class(a,b)时将使用__call__方法？我的理解是，在调用实例时使用__call__方法，而我还没有实例。 - ibrewster

添加 __call__ 不是你想要的。这将让你使用 x = my_func(); x()，但在这种情况下并没有帮助。 - David Wolever

Answer 1

您将会面临的困难是：为了使一个函数既可以作为上下文管理器（with foo() as x），又可以作为普通函数（x = foo()）使用，从函数返回的对象需要同时拥有__enter__和__exit__方法，而在一般情况下，没有很好的方法来向现有对象添加方法。

一种方法可能是创建一个包装类，使用__getattr__将方法和属性传递给原始对象：

class ContextWrapper(object):
    def __init__(self, obj):
        self.__obj = obj

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        ... handle __exit__ ...

    def __getattr__(self, attr):
        return getattr(self.__obj, attr)

但是这样会导致一些微妙的问题，因为它与原始函数返回的对象不完全相同（例如，isinstance测试将失败，一些内置函数如iter(obj)也不能正常工作等）。

您还可以动态地对返回的对象进行子类化，如此处所示：https://dev59.com/vXM_5IYBdhLWcg3wSBAU#1445289：

class ObjectWrapper(BaseClass):
    def __init__(self, obj):
        self.__class__ = type(
            obj.__class__.__name__,
            (self.__class__, obj.__class__),
            {},
        )
        self.__dict__ = obj.__dict__

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        ... handle __exit__ ...

但是这种方法也存在问题（如链接的帖子中所述），而且这是一种我个人不太舒服引入的魔法水平，除非有充分的理由。

我通常更喜欢添加显式的__enter__和__exit__方法，或者使用像contextlib.closing这样的辅助工具：

with closing(my_func()) as my_obj:
    … do stuff …

- David Wolever

---

啊，我缺少的关键是 返回的 对象需要 __enter__ 和 __exit__ 函数 - 不一定是被调用的对象（在这种情况下是函数）需要它们。所以通过将这些函数添加到由函数返回的类中，并使 __enter__ 函数简单地返回 self，就可以让它工作了！ - ibrewster

Answer 2

为了更加清晰：如果您能够更改my_class，那么您当然会向该类添加__enter__/__exit__描述符。

如果您不能更改my_class（根据您的问题所推断），那么这就是我所提到的解决方案：

class my_class(object):

    def __init__(self, result):
        print("this works", result)

class manage_me(object):

    def __init__(self, callback):
        self.callback = callback

    def __enter__(self):
        return self

    def __exit__(self, ex_typ, ex_val, traceback):
        return True

    def __call__(self, *args, **kwargs):
        return self.callback(*args, **kwargs)


def my_func(arg_1,arg_2):
    result=arg_1+arg_2
    return my_class(result)


my_func_object = manage_me(my_func) 

my_func_object(1, 1)
with my_func_object as mf:
    mf(1, 2)

作为一个装饰器：

@manage_me
def my_decorated_func(arg_1, arg_2):
    result = arg_1 + arg_2
    return my_class(result)

my_decorated_func(1, 3)
with my_decorated_func as mf:
    mf(1, 4)

- apex-meme-lord

---

我可以轻松地更改my_class。我的困惑源于我不直接调用（或实例化）my_class——相反，我调用一个返回my_class实例的函数，并且我希望能够像open()函数一样把该函数作为函数或上下文管理器使用。我没有意识到需要将返回的对象表现为上下文管理器 - 我以为我所调用的函数需要是一个上下文管理器。 - ibrewster