How to accept both filenames and file-like objects in Python functions?
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.
精彩评论