Examples of Python context managers in action

context_manager_examples.py

import sys
from io import StringIO

"""
Python Context Managers: Some Examples
This file contains a few examples of ways you can
use context managers. These are minimal examples that
should generalize usefully in many situations.

Note: a context manager is opened like this:

with MyContextObject(*args,**kwargs) as var:
    ...
"""


# Example:
# This context manager replaces stdout with a StringIO buffer,
# and on exiting returns the list of output lines in the buffer
# (and restores stdout).
class CaptureStdout(list):
    """An object to help capture stdout from Snakemake"""
    def __init__(self,passthru=False):
        # Boolean: should we pass everything through to stdout?
        # (this object is only functional if passthru is False)
        self.passthru = passthru

    def __enter__(self):
        # Open a new context with this CaptureStdout
        # object. This happens when we say
        # "with CatpureStdout('label') as output:"

        # If we are just passing input on to output, pass thru
        if self.passthru:
            return self

        # Otherwise, we want to swap out sys.stdout with
        # a StringIO object that will save stdout.
        #
        # Save the existing stdout object so we can
        # restore it when we're done
        self._stdout = sys.stdout
        # Now swap out stdout
        sys.stdout = self._stringio = StringIO()
        return self

    def __exit__(self, *args):
        # If we are just passing input on to output, pass thru
        if self.passthru:
            return self

        # This entire class extends the list class,
        # so we call self.extend() to add a list to
        # the end of self (in this case, all the new
        # lines from our StringIO object).
        self.extend(self._stringio.getvalue().splitlines())
        del self._stringio    # free up some memory

        # Clean up by setting sys.stdout back to what
        # it was before we opened up this context.
        sys.stdout = self._stdout

# To use the CaptureStdout context manager,
# set the passthru boolean to whether you
# want to pass output on to stdout (passthru=True)
# or whether you want to capture output to
# the output variable (passthru=False)
with CaptureStdout(passthru=False) as output:
  for i in range(10):
    print("Hello world %d"%(i+1))

# Now we can do something with output
with open('output.file','w') as f:
  f.write("\n".join(output))