Use a context manager for python script output to a file?

Question

I'm programming a script where I have an option, to be passed on the command line, whether the script should print its results to stdout or to a predefined results file. A code outline for this is shown below. I now have read a little bit about context managers in Python, but am not really sure whether and how to use a context manager in this specific situation. So I am looking for advice

1. whether it makes sense to use a context manager in this problem
2. how to go about implementing it.

So, the code without context manager is:

option_file = True # would come from OptionParser in real code

if option_file:
    out = open("resultsfile", "w")
else:
    out = sys.stdout

# do some computations
out.write("Results of script")
# more computations and calls to out.write

if option_file:
    out.close()

Answer 1

A context manager is something you can use with a with statement. It is explicitly designed to:

• perform some setup,
• give you an object, and
• perform some teardown again (even if you raise an exception).

For example, open can be used as a context manager. In the following code

with open(...) as f:
    # do stuff

it doesn't matter what stuff is, the file will always be closed. (Well, usually. Except in silly cases like the power being turned off or the process being killed.)

You should use a context manager when you are in this sort of situation. It doesn't look to me like you are, so I see no reason for a context manager.

---

There is an alternative (not better or worse, just different) way of writing your code that does use a context manager. If you want to redirect stdout temporarily -- but ensure that you restore it when you're done -- then you are in the situation above. Here's an example:

@contextlib.contextmanager
def redirect_stdout(stream):
    import sys
    sys.stdout = stream
    yield
    sys.stdout = sys.__stdout__

You can then write code like

with open(...) as f:
    with redirect_stdout(f):
        # do stuff

and any writes to stdout in stuff will go to f instead.

---

EDIT: You are correct that there is no way conditionally to have a context manager: either you are in one or you aren't. You could always write your own which might do nothing:

@contextlib.contextmanager
def maybe_open(path, do_nothing=True):
    if do_nothing:
        f = None
        yield sys.stdout
    else:
        f = open(path)
        yield f

    if f:
        f.close()

This is almost certainly overkill.