如何在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

然而,这两种解决方案都没有让我太过信服,尤其是第二种解决方案,因为我看到了太多的陷阱.

什么是最简洁(和最Pythonic)的方式接受文件类对象或字符串到文本读取功能？

Answer 1

不要接受文件和字符串.如果你打算接受类文件对象,那么就意味着你将不会检查类型,只需调用实际参数(所需的方法read,write等等).如果您要接受字符串,那么您将结束open文件,这意味着您将无法模拟参数.所以我会说接受文件,让调用者传递一个类文件对象,不要检查类型.

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

我正在使用上下文管理器包装器。当它是文件名 (str) 时，退出时关闭文件。

@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