15c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# Copyright (C) 2012 Google, Inc.
25c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# Copyright (C) 2010 Chris Jerdonek (cjerdonek@webkit.org)
35c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)#
45c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# Redistribution and use in source and binary forms, with or without
55c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# modification, are permitted provided that the following conditions
65c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# are met:
75c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# 1.  Redistributions of source code must retain the above copyright
85c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)#     notice, this list of conditions and the following disclaimer.
95c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# 2.  Redistributions in binary form must reproduce the above copyright
105c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)#     notice, this list of conditions and the following disclaimer in the
115c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)#     documentation and/or other materials provided with the distribution.
125c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)#
135c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# THIS SOFTWARE IS PROVIDED BY APPLE INC. AND ITS CONTRIBUTORS ``AS IS'' AND
145c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
155c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
165c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# DISCLAIMED. IN NO EVENT SHALL APPLE INC. OR ITS CONTRIBUTORS BE LIABLE FOR
175c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
185c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
195c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
205c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
215c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
225c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
235c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
245c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)import logging
255c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)import StringIO
265c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
27926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)from webkitpy.common.system.systemhost import SystemHost
285c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)from webkitpy.layout_tests.views.metered_stream import MeteredStream
295c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
305c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)_log = logging.getLogger(__name__)
315c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
325c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
335c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)class Printer(object):
345c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)    def __init__(self, stream, options=None):
355c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.stream = stream
365c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.meter = None
375c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.options = options
385c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.num_tests = 0
395c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.num_completed = 0
405c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.num_errors = 0
415c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.num_failures = 0
425c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.running_tests = []
435c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.completed_tests = []
445c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        if options:
455c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            self.configure(options)
465c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
475c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)    def configure(self, options):
485c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.options = options
495c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
505c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        if options.timing:
515c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            # --timing implies --verbose
525c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            options.verbose = max(options.verbose, 1)
535c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
545c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        log_level = logging.INFO
555c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        if options.quiet:
565c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            log_level = logging.WARNING
575c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        elif options.verbose == 2:
585c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            log_level = logging.DEBUG
595c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
60926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)        self.meter = MeteredStream(self.stream, (options.verbose == 2),
61926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)            number_of_columns=SystemHost().platform.terminal_width())
625c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
635c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        handler = logging.StreamHandler(self.stream)
645c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # We constrain the level on the handler rather than on the root
655c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # logger itself.  This is probably better because the handler is
665c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # configured and known only to this module, whereas the root logger
675c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # is an object shared (and potentially modified) by many modules.
685c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # Modifying the handler, then, is less intrusive and less likely to
695c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # interfere with modifications made by other modules (e.g. in unit
705c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # tests).
715c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        handler.name = __name__
725c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        handler.setLevel(log_level)
735c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        formatter = logging.Formatter("%(message)s")
745c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        handler.setFormatter(formatter)
755c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
765c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        logger = logging.getLogger()
775c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        logger.addHandler(handler)
785c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        logger.setLevel(logging.NOTSET)
795c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
805c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # Filter out most webkitpy messages.
815c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        #
825c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # Messages can be selectively re-enabled for this script by updating
835c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # this method accordingly.
845c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        def filter_records(record):
857757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            """Filter out non-third-party webkitpy messages."""
865c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            # FIXME: Figure out a way not to use strings here, for example by
875c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            #        using syntax like webkitpy.test.__name__.  We want to be
885c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            #        sure not to import any non-Python 2.4 code, though, until
895c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            #        after the version-checking code has executed.
907757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            if (record.name.startswith("webkitpy.test")):
915c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)                return True
925c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            if record.name.startswith("webkitpy"):
935c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)                return False
945c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            return True
955c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
965c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        testing_filter = logging.Filter()
975c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        testing_filter.filter = filter_records
985c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
995c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # Display a message so developers are not mystified as to why
1005c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # logging does not work in the unit tests.
1015c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        _log.info("Suppressing most webkitpy logging while running unit tests.")
1025c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        handler.addFilter(testing_filter)
1035c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1045c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        if self.options.pass_through:
105926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)            # FIXME: Can't import at top of file, as outputcapture needs unittest2
106926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)            from webkitpy.common.system import outputcapture
1075c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            outputcapture.OutputCapture.stream_wrapper = _CaptureAndPassThroughStream
1085c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1095c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)    def write_update(self, msg):
1105c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.meter.write_update(msg)
1115c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1125c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)    def print_started_test(self, source, test_name):
1135c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.running_tests.append(test_name)
1145c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        if len(self.running_tests) > 1:
1155c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            suffix = ' (+%d)' % (len(self.running_tests) - 1)
1165c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        else:
1175c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            suffix = ''
1185c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1195c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        if self.options.verbose:
1205c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            write = self.meter.write_update
1215c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        else:
1225c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            write = self.meter.write_throttled_update
1235c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1245c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        write(self._test_line(self.running_tests[0], suffix))
1255c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1265c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)    def print_finished_test(self, source, test_name, test_time, failures, errors):
1275c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        write = self.meter.writeln
1285c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        if failures:
1295c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            lines = failures[0].splitlines() + ['']
1305c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            suffix = ' failed:'
1315c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            self.num_failures += 1
1325c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        elif errors:
1335c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            lines = errors[0].splitlines() + ['']
1345c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            suffix = ' erred:'
1355c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            self.num_errors += 1
1365c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        else:
1375c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            suffix = ' passed'
1385c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            lines = []
1395c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            if self.options.verbose:
1405c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)                write = self.meter.writeln
1415c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            else:
1425c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)                write = self.meter.write_throttled_update
1435c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        if self.options.timing:
1445c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            suffix += ' %.4fs' % test_time
1455c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1465c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.num_completed += 1
1475c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1485c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        if test_name == self.running_tests[0]:
1495c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            self.completed_tests.insert(0, [test_name, suffix, lines])
1505c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        else:
1515c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            self.completed_tests.append([test_name, suffix, lines])
1525c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.running_tests.remove(test_name)
1535c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1545c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        for test_name, msg, lines in self.completed_tests:
1555c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            if lines:
1565c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)                self.meter.writeln(self._test_line(test_name, msg))
1575c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)                for line in lines:
1585c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)                    self.meter.writeln('  ' + line)
1595c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            else:
1605c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)                write(self._test_line(test_name, msg))
1615c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self.completed_tests = []
1625c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1635c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)    def _test_line(self, test_name, suffix):
164926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)        format_string = '[%d/%d] %s%s'
165926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)        status_line = format_string % (self.num_completed, self.num_tests, test_name, suffix)
166926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)        if len(status_line) > self.meter.number_of_columns():
167926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)            overflow_columns = len(status_line) - self.meter.number_of_columns()
168926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)            ellipsis = '...'
169926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)            if len(test_name) < overflow_columns + len(ellipsis) + 3:
170926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)                # We don't have enough space even if we elide, just show the test method name.
171926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)                test_name = test_name.split('.')[-1]
172926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)            else:
173926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)                new_length = len(test_name) - overflow_columns - len(ellipsis)
174926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)                prefix = int(new_length / 2)
175926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)                test_name = test_name[:prefix] + ellipsis + test_name[-(new_length - prefix):]
176926b001d589ce2f10facb93dd4b87578ea35a855Torne (Richard Coles)        return format_string % (self.num_completed, self.num_tests, test_name, suffix)
1775c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1785c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)    def print_result(self, run_time):
1795c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        write = self.meter.writeln
1805c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        write('Ran %d test%s in %.3fs' % (self.num_completed, self.num_completed != 1 and "s" or "", run_time))
1815c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        if self.num_failures or self.num_errors:
1825c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            write('FAILED (failures=%d, errors=%d)\n' % (self.num_failures, self.num_errors))
1835c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        else:
1845c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            write('\nOK\n')
1855c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1865c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1875c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)class _CaptureAndPassThroughStream(object):
1885c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)    def __init__(self, stream):
1895c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self._buffer = StringIO.StringIO()
1905c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self._stream = stream
1915c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1925c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)    def write(self, msg):
1935c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self._stream.write(msg)
1945c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
1955c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # Note that we don't want to capture any output generated by the debugger
1965c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # because that could cause the results of capture_output() to be invalid.
1975c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        if not self._message_is_from_pdb():
1985c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)            self._buffer.write(msg)
1995c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
2005c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)    def _message_is_from_pdb(self):
2015c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # We will assume that if the pdb module is in the stack then the output
2025c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # is being generated by the python debugger (or the user calling something
2035c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        # from inside the debugger).
2045c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        import inspect
2055c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        import pdb
2065c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        stack = inspect.stack()
2075c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        return any(frame[1] == pdb.__file__.replace('.pyc', '.py') for frame in stack)
2085c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
2095c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)    def flush(self):
2105c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        self._stream.flush()
2115c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)
2125c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)    def getvalue(self):
2135c87bf8b86a7c82ef50fb7a89697d8e02e2553beTorne (Richard Coles)        return self._buffer.getvalue()
214