Fix Python – How to handle both `with open(…)` and `sys.stdout` nicely?

Question

Asked By – Jakub M.

Often I need to output data either to file or, if file is not specified, to stdout. I use the following snippet:

if target:
    with open(target, 'w') as h:
        h.write(content)
else:
    sys.stdout.write(content)

I would like to rewrite it and handle both targets uniformly.

In ideal case it would be:

with open(target, 'w') as h:
    h.write(content)

but this will not work well because sys.stdout is be closed when leaving with block and I don’t want that. I neither want to

stdout = open(target, 'w')
...

because I would need to remember to restore original stdout.

Related:

Edit

I know that I can wrap target, define separate function or use context manager. I look for a simple, elegant, idiomatic solution fitting that wouldn’t require more than 5 lines

Now we will see solution for issue: How to handle both `with open(…)` and `sys.stdout` nicely?


Answer

Just thinking outside of the box here, how about a custom open() method?

import sys
import contextlib

@contextlib.contextmanager
def smart_open(filename=None):
    if filename and filename != '-':
        fh = open(filename, 'w')
    else:
        fh = sys.stdout

    try:
        yield fh
    finally:
        if fh is not sys.stdout:
            fh.close()

Use it like this:

# For Python 2 you need this line
from __future__ import print_function

# writes to some_file
with smart_open('some_file') as fh:
    print('some output', file=fh)

# writes to stdout
with smart_open() as fh:
    print('some output', file=fh)

# writes to stdout
with smart_open('-') as fh:
    print('some output', file=fh)

This question is answered By – Wolph

This answer is collected from stackoverflow and reviewed by FixPython community admins, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0