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)

Answer 2

It's trivial to create one with the @contextlib.contextmanager decorator:

from contextlib import contextmanager

@contextmanager
def dummy_manager(ob):
    yield ob

That's it; this creates a context manager that does nothing but hand you ob back, and the __exit__ handler does exactly nothing.

I'd use it like this:

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

Answer 3

You don't need a dummy manager in your case. sys.stdin, being file-like, can be used as a context manager.

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

One thing to note is that when __exit__ing the block, sys.stdin gets closed (while you ordinarily don't need/want to close it yourself), but that shouldn't be a problem.

Answer 4

Various ArgParse wrappers and replacements now exist which support this nicely. I like click:

with click.open_file(filename) as lines:
    for line in lines:
        process(line)

This will handle sys.stdin if filename is - and otherwise fall back to the regular open with a close hidden in the context manager's finally part.