打开一个上下文管理器，但前提是调用者没有提供

Question

这是一个人为的例子，我的实际用例是使用（主要是只读的）数据库连接 - 如果未提供连接，请打开一个。

我已将其更改为由with open(xxx) as fo 提供的打开文件句柄。这个想法是我有一些可能需要做一些工作的函数——如果他们的调用者已经给了他们一个上下文管理的对象（在这种情况下是文件句柄），请使用它。否则用它创建一个本地上下文管理器。

从概念上讲，这就是我想要的，但它按预期失败了：

def write_it(name, fo=None):
    if not fo:
       with open(name,"w") as fo:
            #I'd want to keep `fo` but it wont work
            pass

    #assume this is a lot of complex code

    #if fo was opened here, it will have been closed already
    fo.write(name)


name = "works_w_context.txt"
with open(name, "w") as fo:
    write_it(name, fo)
    fo.write("\n and it still should be open")

#this fails, as expected
name = "error_wo_context.txt"
write_it(name)

不出所料，当我没有提供打开的文件时出现错误。

Traceback (most recent call last):
  File "test_182_context.py", line 21, in <module>
    write_it(name)
  File "test_182_context.py", line 12, in write_it
    fo.write(data)
ValueError: I/O operation on closed file.

(venv38) myuser@test_182_context$ dir *.txt
-rw-r--r--  1 myuser  staff  52 Oct  1 17:26 works_w_context.txt
-rw-r--r--  1 myuser  staff   0 Oct  1 17:26 error_wo_context.txt

我的解决方法 - 有没有更好的方法使用 contextlib 模块？

我发现这样做的唯一方法是创建一个公共存根函数，该函数在需要时打开文件，然后使用文件句柄调用实际的函数。

但正如你所见，我现在已经重复了函数签名 4 次。当然，*args 和 **kwargs 可以提供一点帮助，但它们也会使代码更难遵循。

def _write_it(name, fo=None):
    #assume this is a lot of complex code
    data = "some very complicated calculations taking many lines"    
    fo.write(data)


def write_it(name, fo=None):
    if not fo:
       with open(name,"w") as fo:
            _write_it(name, fo)
    else:
        _write_it(name, fo)



name = "works2_w_context.txt"
with open(name, "w") as fo:
    write_it(name, fo)
    fo.write("\n and it still should be open")

#this fails, as expected
name = "works2_wo_context.txt"
write_it(name)

是的，它有效：

(venv38) myuser@test_182_context$ py test_182_context_2.py
(venv38) myuser@test_182_context$ dir *.txt
-rw-r--r--  1 myuser  staff  52 Oct  1 17:29 works2_w_context.txt
-rw-r--r--  1 myuser  staff  52 Oct  1 17:29 works2_wo_context.txt

今天早些时候有人问了一个问题并提到了contextlib.nullcontext的评论

返回一个从 enter 返回 enter_result 的上下文管理器，否则什么也不做。它旨在用作可选上下文管理器的替代品。

它有点看起来与我所追求的有关，但同时它看起来并不能解决嵌套的核心问题。

Answer 1

您可以使用退出堆栈有条件地为您自己打开的文件添加新的上下文管理器。

from contextlib import ExitStack

def write_it(name, fo=None):
    with ExitStack() as es:
        if fo is None:
            fo = es.enter_context(open(name, "w"))
    
        data = ...
        fo.write()

如果退出堆栈为空，则在with 语句退出时没有任何事情可做。只会将新打开的文件添加到堆栈中，并在 with 语句完成时关闭。

---

如果fo 不是None，您也可以使用空上下文管理器：

from contextlib import nullcontext

def write_it(name, fo=None):
    if fo is None:
        # With no file received, open a new one and use it
        # as the context manager
        cm = open(name, "w")
    else:
        # Create a do-nothing context manager that won't close
        # the received file.
        cm = nullcontext(fo)

    # Get an open file handle from the defined context manager
    with cm as f:
        data = ...
        f.write(data)

空上下文管理器的__enter__ 方法返回包装的对象，但它的__exit__ 方法什么也不做。