如何在Python函数中同时接受文件名和类文件对象？

Question

在我的代码中，有一个叫做 load_dataset 的函数，它用来读取一个文本文件并进行一些处理。最近我在想，能不能让这个函数也支持类似文件的对象，于是我开始考虑最好的实现方法。目前我有两个想法：

第一个是类型检查：

if isinstance(inputelement, basestring):
   # open file, processing etc
# or
# elif hasattr(inputelement, "read"):
elif isinstance(inputelement, file):
   # Do something else

另外一个是使用两个不同的参数：

def load_dataset(filename=None, stream=None):
    if filename is not None and stream is None:
        # open file etc
    elif stream is not None and filename is None:
        # do something else

不过这两种方案都让我不太满意，尤其是第二种，因为我觉得会有很多潜在的问题。

那么，接受一个类似文件的对象或者字符串的函数，最简单（而且最符合Python风格）的做法是什么呢？

Answer 1

我在使用一个上下文管理器的包装器。当传入的是文件名（字符串）时，退出时会自动关闭这个文件。

@contextmanager
def fopen(filein, *args, **kwargs):
    if isinstance(filein, str):  # filename
        with open(filein, *args, **kwargs) as f:
            yield f
    else:  # file-like object
        yield filein

然后你可以像这样使用它：

with fopen(filename_or_fileobj) as f:
    # do sth. with f

Answer 2

一种可以让你传入文件名或者类似文件的对象作为参数的方法，就是实现一个上下文管理器，它能够同时处理这两种情况。你可以在这里找到具体的实现，我引用这段内容是为了提供一个完整的答案：

class open_filename(object):
"""Context manager that opens a filename and closes it on exit, but does
nothing for file-like objects.
"""
def __init__(self, filename, *args, **kwargs):
    self.closing = kwargs.pop('closing', False)
    if isinstance(filename, basestring):
        self.fh = open(filename, *args, **kwargs)
        self.closing = True
    else:
        self.fh = filename

def __enter__(self):
    return self.fh

def __exit__(self, exc_type, exc_val, exc_tb):
    if self.closing:
        self.fh.close()

    return False

接下来可能的用法是：

def load_dataset(file_):
    with open_filename(file_, "r") as f:
        # process here, read only if the file_ is a string

Answer 3

不要同时接受文件和字符串。如果你要接受像文件一样的对象，那就意味着你不会检查类型，而是直接调用实际参数需要的方法（比如 read、write 等）。如果你要接受字符串，那你就得去 open 文件，这样就没办法模拟这些参数。所以我建议你只接受文件，让调用者传给你一个像文件的对象，不用检查类型。