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)