开发者

How to accept both filenames and file-like objects in Python functions?

开发者 https://www.devze.com 2023-04-01 00:57 出处:网络
In my code, I have a load_dataset function that reads a text file and does some processing. Recently I thought about adding support to file-like objects, and I wondered over the best approach to this.

In my code, I have a load_dataset function that reads a text file and does some processing. Recently I thought about adding support to file-like objects, and I wondered over the best approach to this. Currently I have two implementations in mind:

First, type checking:

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

Alternatively, two different arguments:

def load_dataset(filename=None, stream=None):
    if filename is not None and stream is None:
        开发者_JS百科# open file etc
    elif stream is not None and filename is None:
        # do something else

Both solutions however don't convince me too much, especially the second as I see way too many pitfalls.

What is the cleanest (and most Pythonic) way to accept a file-like object or string to a function that does text reading?


One way of having either a file name or a file-like object as argument is the implementation of a context manager that can handle both. An implementation can be found here, I quote for the sake of a self contained answer:

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

Possible usage then:

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


Don't accept both files and strings. If you're going to accept file-like objects, then it means you won't check the type, just call the required methods on the actual parameter (read, write, etc.). If you're going to accept strings, then you're going to end up open-ing files, which means you won't be able to mock the parameters. So I'd say accept files, let the caller pass you a file-like object, and don't check the type.


I'm using a context manager wrapper. When it's a filename (str), close the file on exit.

@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

Then you can use it like:

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


Python follows duck typing, you may check that this is file object by function that you need from object. For example hasattr(obj, 'read') against isinstance(inputelement, file). For converting string to file objects you may also use such construction:

if not hasattr(obj, 'read'):
    obj = StringIO(str(obj))

After this code you will safely be able to use obj as a file.

0

精彩评论

暂无评论...
验证码 换一张
取 消