Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Pythonic way to use context manager conditionally

This is something I think must come up quite often but I haven't been able to find a good solution for it. Say I have a function which may be passed an open resource as an argument (like a file or database connection object) or needs to create one itself. If the function needs to open a file on its own, best practice is usually considered something like:

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

to ensure the file is always closed when the with block is exited. However if an existing file handle is passed in the function should probably not close it itself.

Consider the following function which takes either an open file object or a string giving a path to the file as its argument. If it is passed a file path it should probably be written as above. Otherwise the with statement should be omitted. This results in duplicate code:

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...

This could of course be avoided by defining a helper function and calling it in both places, but this seems inelegant. A better way I thought of was to define a context manager class that wraps an object like so:

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...

This works but unless there's a built-in object that does this (I don't think there is) I either have to copy-paste that object everywhere or always have to import my custom utilities module. I'm feeling like there's a simpler way to do this that I've missed.

like image 210
JaredL Avatar asked Nov 09 '22 00:11

JaredL


1 Answers

However, I prefer, I don't know how pythonic it is, but it's straightforward

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

the problem could be solved nicer with singledispatch from python 3.4

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
like image 162
kwarunek Avatar answered Nov 15 '22 06:11

kwarunek