Fallback to stdout if no file name provided

Question

I have a script that accepts as an argument a filename than opens it and writes some stuff.

I use the with statement:

with open(file_name, 'w') as out_file:
    ...
    out_file.write(...)

Now what if I want to write to sys.stdout if no file_name is provided?

Do I necessarily need to wrap all actions in a function and put a condition before?

if file_name is None:
    do_everything(sys.stdout)
else:
    with open(file_name, 'w') as out_file:
        do_everything(out_file)

Answer 1

from contextlib import contextmanager
@contextmanager
def file_or_stdout(file_name):
    if file_name is None:
        yield sys.stdout
    else:
        with open(file_name, 'w') as out_file:
            yield out_file

Then you can do

with file_or_stdout(file_name) as wfile:
    do_stuff_writing_to(wfile)

Answer 2

How do you handle command line arguments? If you use argparse you could use the type and default parameters of add_argument to handle this. For example, try something like the following:

import sys
import argparse

def main(argv=None):

    if argv is None:
        argv=sys.argv[1:]

    parser = argparse.ArgumentParser()

    parser.add_argument('infile', nargs='?',
        type=argparse.FileType('w'),
        default=sys.stdin)

    args = parser.parse_args(argv)

    print args.infile

    return 0

if __name__=="__main__":
    sys.exit(main(sys.argv[1:]))

If a file name is present as an argument the the script argparse will automatically open and close this file and args.infile will be a handle to this file. Otherwise args.infile will simply be sys.stdin.