CODEWITHSUNDEEP
pythonpython-2.7contextmanager
JaredL
JaredL

Reputation: 1443

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.

Upvotes: 10

Views: 3202

Answers (4)

Clay
Clay

Reputation: 65

How about recursion?

def foo(f):
    if isinstance(f, basestring):
        # Path to file, need to open
        with open(f) as fh:
            # recurse with now-opened file
            return foo(fh)
    # Assume open file
    # do stuff

Call the function, if you have a path rather than the desired object, just create the desired object, pass it as an argument to the same function and return the output.

It feels nice and elegant for a case like your example but the downside is it may not work if you need to do it within a loop or something.

Upvotes: 0

Kache
Kache

Reputation: 16677

That "context manager wrapper" idea is the right way to go. Not only is it lighterweight written as a function:

from contextlib import contextmanager

@contextmanager
def nullcontext(obj=None):
    yield obj

A sign that it's the right choice: the purpose-built nullcontext is available in the stdlib as of 3.7 (and not at time of asking/for python-2.7), with your exact use-case in the docs:

def process_file(file_or_path):
    if isinstance(file_or_path, str):
        # If string, open file
        cm = open(file_or_path)
    else:
        # Caller is responsible for closing file
        cm = nullcontext(file_or_path)

    with cm as file:
        # Perform processing on the file

Upvotes: 4

eestrada
eestrada

Reputation: 1603

This solution avoids an explicit boolean value like f_own (that is mentioned in a comment of the answer from @kAlmAcetA), and instead just checks the identity of the input parameter f to the file handle fh. A try/finally clause is the only way to do this without creating a helper class as a context manager.

def foo(f):
    fh = open(f) if isinstance(f, basestring) else f

    try:
        # do stuff...
    finally:
        if fh is not f:
            fh.close()

If you need to do something like this in more than one function then, yes, you probably should create a utility module with a context manager class to do it, like this:

class ContextWrapper(object):
    def __init__(self, file):
        self.f = file

    def __enter__(self):
        self.fh = open(self.f) if isinstance(self.f, basestring) else self.f
        return self.fh

    def __exit__(self, *args):
        if self.fh is not self.f:
            self.fh.close()

Then you could unconditionally wrap like this:

def foo(f):
    with ContextManager(f) as fh:
        # do stuff...

Upvotes: 0

kwarunek
kwarunek

Reputation: 12577

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

Upvotes: 1

Related Questions