Context manager for optionally redirected I/O

Question

I'm often faced with the situation that depending on some command line argument, input might either come from a file or standard input. The same goes for output. I really like how context managers in python 3 work, and therefore try to make all my open calls part of some with statement. But in this case, I'm having trouble.

if args.infile:
    with open(args.infile, "r") as f:
        process(f)
else:
    process(sys.stdin)

is already clumsy, and with both input and output I'd have to cater for four combinations. I would like something easier, e.g.

with (open(args.infile, "r") if args.infile
      else DummyManager(sys.stdin)) as f:
    process(f)

Is there something like this DummyManager in the python standard libraries? Something which implements the context manager protocol, but only to return a fixed value from its __enter__ method? I guess the most likely location for such a class would be contextlib, and since I didn't find anything like this there, perhaps there is no such thing. Are there other elegant solutions you can suggest?

Answer 1

In your case, you could use fileinput module:

from fileinput import FileInput

with FileInput(args.infile) as file:
    process(file)
# sys.stdin is still open here

If args.infile='-' then it uses sys.stdin. If inplace=True parameter then it redirects sys.stdout to the input file. You could pass several filenames. If there are no filenames it uses filenames given at the command-line or stdin.

Or you could leave the file as is:

import sys
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--log', default=sys.stdout, type=argparse.FileType('w'))
args = parser.parse_args()
with args.log:
    args.log.write('log message')
# sys.stdout may be closed here

It should be fine for most programs where stdout may be used to write the result.

To avoid closing sys.stdin / sys.stdout, you could use ExitStack to enable the context managers conditionally:

from contextlib import ExitStack

with ExitStack() as stack:
    if not args.files:
       files = [sys.stdin]
    else:
       files = [stack.enter_context(open(name)) for name in args.files]

    if not args.output:
       output_file = sys.stdout
       stack.callback(output_file.flush) # flush instead of closing 
    else:
       output_file = stack.enter_context(open(args.output, 'w'))

    process(files, output_file)