Class AssertPrints

IPython/testing/tools.py:339–379

Context manager for testing that code prints certain text. Examples -------- >>> with AssertPrints("abc", suppress=False): ... print("abcd") ... print("def") ... abcd def

"""

class AssertPrints(object):
    """Context manager for testing that code prints certain text.

    Examples
    --------
    >>> with AssertPrints("abc", suppress=False):
    ...     print("abcd")
    ...     print("def")
    ...
    abcd
    def
    """
    def __init__(self, s, channel='stdout', suppress=True):
        self.s = s
        if isinstance(self.s, (str, _re_type)):
            self.s = [self.s]
        self.channel = channel
        self.suppress = suppress

    def __enter__(self):
        self.orig_stream = getattr(sys, self.channel)
        self.buffer = MyStringIO()
        self.tee = Tee(self.buffer, channel=self.channel)
        setattr(sys, self.channel, self.buffer if self.suppress else self.tee)

    def __exit__(self, etype, value, traceback):
        try:
            if value is not None:
                # If an error was raised, don't check anything else
                return False
            self.tee.flush()
            setattr(sys, self.channel, self.orig_stream)
            printed = self.buffer.getvalue()
            for s in self.s:
                if isinstance(s, _re_type):
                    assert s.search(printed), notprinted_msg.format(s.pattern, self.channel, printed)
                else:
                    assert s in printed, notprinted_msg.format(s, self.channel, printed)
            return False
        finally:
            self.tee.close()

printed_msg = """Found {0!r} in printed output (on {1}):
-------