using Python 'with' statement with sys.stdout

Question

I always open and write into files using with statement:

with open('file_path', 'w') as handle:
    print >>handle, my_stuff

However, there is one instance where I need to be able to be more flexible, and write to sys.stdout (or other types of streams), if that is provided instead of file path:

So, my question is this: Is there a way for using with statement both with real files and with sys.stdout?

Note that I can use the following code, but I think this defeats the purpose of using with:

if file_path != None:
    outputHandle = open(file_path, 'w')
else:
    outputHandle = sys.stdout

with outputHandle as handle:
    print >>handle, my_stuff

Answer 1

You can create a context manager and use it like this

import contextlib, sys

@contextlib.contextmanager
def file_writer(file_name = None):
    # Create writer object based on file_name
    writer = open(file_name, "w") if file_name is not None else sys.stdout
    # yield the writer object for the actual use
    yield writer
    # If it is file, then close the writer object
    if file_name != None: writer.close()

with file_writer("Output.txt") as output:
    print >>output, "Welcome"

with file_writer() as output:
    print >>output, "Welcome"

If you don't pass any input to file_writer it will use sys.stdout.

Answer 2

Thing is, you don't need to use a context processor with stdout, because you're not opening or closing it. A less fancy way of abstracting this is:

def do_stuff(file):
    # Your real code goes here. It works both with files or stdout
    return file.readline()

def do_to_stdout():
    return do_stuff(sys.stdout)

def do_to_file(filename):
    with open(filename) as f:
        return do_stuff(f)


print do_to_file(filename) if filename else do_to_stdout()