How to handle both `with open(...)` and `sys.stdout` nicely?

An improvement of Wolph's answer

import sys
import contextlib

@contextlib.contextmanager
def smart_open(filename: str, mode: str = 'r', *args, **kwargs):
    '''Open files and i/o streams transparently.'''
    if filename == '-':
        if 'r' in mode:
            stream = sys.stdin
        else:
            stream = sys.stdout
        if 'b' in mode:
            fh = stream.buffer  # type: IO
        else:
            fh = stream
        close = False
    else:
        fh = open(filename, mode, *args, **kwargs)
        close = True

    try:
        yield fh
    finally:
        if close:
            try:
                fh.close()
            except AttributeError:
                pass

This allows binary IO and pass eventual extraneous arguments to open if filename is indeed a file name.


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)

Stick with your current code. It's simple and you can tell exactly what it's doing just by glancing at it.

Another way would be with an inline if:

handle = open(target, 'w') if target else sys.stdout
handle.write(content)

if handle is not sys.stdout:
    handle.close()

But that isn't much shorter than what you have and it looks arguably worse.

You could also make sys.stdout unclosable, but that doesn't seem too Pythonic:

sys.stdout.close = lambda: None

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

Tags:

Python