有条件地使用上下文管理器的 Pythonic 方式

Question

我认为这是经常出现的问题，但我一直无法找到一个好的解决方案。假设我有一个函数，可以将开放资源作为参数传递（如文件或数据库连接对象），或者需要自己创建一个函数。如果函数需要自己打开文件，最佳实践通常被认为是这样的：

with open(myfile) as fh:
    # do stuff with open file handle...

确保with退出块时文件始终关闭。但是，如果在函数中传递现有文件句柄，则该函数可能不应该关闭它本身。

考虑以下函数，它接受一个打开的文件对象或一个给出文件路径的字符串作为其参数。如果传递的是文件路径，则可能应该按上面的方式编写。否则该with语句应被省略。这会导致重复的代码：

def foo(f):
    if isinstance(f, basestring):
        # Path to file, need to open
        with open(f) as fh:
            # do stuff with fh...
    else:
        # Assume open file
        fh = f
        # do the same stuff...

当然，可以通过定义一个辅助函数并在两个地方调用它来避免这种情况，但这看起来不太优雅。我想到的一个更好的方法是定义一个上下文管理器类来包装一个对象，如下所示：

class ContextWrapper(object):
    def __init__(self, wrapped):
        self.wrapped = wrapped
    def __enter__(self):
        return self.wrapped
    def __exit__(self, *args):
        pass

def foo(f):
    if isinstance(f, basestring):
        cm = open(f)
    else:
        cm = ContextWrapper(f)

    with cm as fh:
        # do stuff with fh...

这是可行的，但除非有一个内置对象可以执行此操作（我认为没有），否则我要么必须将该对象复制粘贴到任何地方，要么总是必须导入我的自定义实用程序模块。我觉得有一种更简单的方法可以做到这一点，但我已经错过了。

Answer 1

不过，我更喜欢，我不知道它有多Pythonic，但它很简单

def foo(f):
    if isinstance(f, basestring):
        f = open(f)
    try:
        # do the stuff
    finally:
        f.close()

使用 python 3.4 中的singledispatch可以更好地解决这个问题

from functools import singledispatch

@singledispatch
def foo(fd):
    with fd as f:
        # do stuff
        print('file')

@foo.register(str)
def _(arg):
    print('string')
    f = open(arg)
    foo(f)


foo('/tmp/file1')  # at first calls registered func and then foo
foo(open('/tmp/file2', 'r'))  # calls foo