diff --git a/gclient_utils.py b/gclient_utils.py index 3a6bc1fa8..dde73def6 100644 --- a/gclient_utils.py +++ b/gclient_utils.py @@ -363,16 +363,21 @@ class Annotated(Wrapper): self.lock = threading.Lock() self.__output_buffers = {} self.__include_zero = include_zero + self._wrapped_write = getattr(self._wrapped, 'buffer', self._wrapped).write @property def annotated(self): return self def write(self, out): + # Store as bytes to ensure Unicode characters get output correctly. + if not isinstance(out, bytes): + out = out.encode('utf-8') + index = getattr(threading.currentThread(), 'index', 0) if not index and not self.__include_zero: # Unindexed threads aren't buffered. - return self._wrapped.write(out) + return self._wrapped_write(out) self.lock.acquire() try: @@ -380,7 +385,7 @@ class Annotated(Wrapper): # Strings are immutable, requiring to keep a lock for the whole dictionary # otherwise. Using an array is faster than using a dummy object. if not index in self.__output_buffers: - obj = self.__output_buffers[index] = [''] + obj = self.__output_buffers[index] = [b''] else: obj = self.__output_buffers[index] finally: @@ -390,18 +395,18 @@ class Annotated(Wrapper): obj[0] += out while True: # TODO(agable): find both of these with a single pass. - cr_loc = obj[0].find('\r') - lf_loc = obj[0].find('\n') + cr_loc = obj[0].find(b'\r') + lf_loc = obj[0].find(b'\n') if cr_loc == lf_loc == -1: break elif cr_loc == -1 or (lf_loc >= 0 and lf_loc < cr_loc): - line, remaining = obj[0].split('\n', 1) + line, remaining = obj[0].split(b'\n', 1) if line: - self._wrapped.write('%d>%s\n' % (index, line)) + self._wrapped_write(b'%d>%s\n' % (index, line)) elif lf_loc == -1 or (cr_loc >= 0 and cr_loc < lf_loc): - line, remaining = obj[0].split('\r', 1) + line, remaining = obj[0].split(b'\r', 1) if line: - self._wrapped.write('%d>%s\r' % (index, line)) + self._wrapped_write(b'%d>%s\r' % (index, line)) obj[0] = remaining def flush(self): @@ -423,7 +428,7 @@ class Annotated(Wrapper): # Don't keep the lock while writting. Will append \n when it shouldn't. for orphan in orphans: if orphan[1]: - self._wrapped.write('%d>%s\n' % (orphan[0], orphan[1])) + self._wrapped_write(b'%d>%s\n' % (orphan[0], orphan[1])) return self._wrapped.flush() diff --git a/tests/gclient_test.py b/tests/gclient_test.py index 1f959ce3e..0faeac44b 100755 --- a/tests/gclient_test.py +++ b/tests/gclient_test.py @@ -76,14 +76,10 @@ class GclientTest(trial_dir.TestCase): self._old_createscm = gclient.gclient_scm.GitWrapper gclient.gclient_scm.GitWrapper = SCMMock SCMMock.unit_test = self - self._old_sys_stdout = sys.stdout - sys.stdout = gclient.gclient_utils.MakeFileAutoFlush(sys.stdout) - sys.stdout = gclient.gclient_utils.MakeFileAnnotated(sys.stdout) def tearDown(self): self.assertEqual([], self._get_processed()) gclient.gclient_scm.GitWrapper = self._old_createscm - sys.stdout = self._old_sys_stdout os.chdir(self.previous_dir) super(GclientTest, self).tearDown() @@ -1366,9 +1362,9 @@ class GclientTest(trial_dir.TestCase): if __name__ == '__main__': sys.stdout = gclient_utils.MakeFileAutoFlush(sys.stdout) - sys.stdout = gclient_utils.MakeFileAnnotated(sys.stdout, include_zero=True) + sys.stdout = gclient_utils.MakeFileAnnotated(sys.stdout) sys.stderr = gclient_utils.MakeFileAutoFlush(sys.stderr) - sys.stderr = gclient_utils.MakeFileAnnotated(sys.stderr, include_zero=True) + sys.stderr = gclient_utils.MakeFileAnnotated(sys.stderr) logging.basicConfig( level=[logging.ERROR, logging.WARNING, logging.INFO, logging.DEBUG][ min(sys.argv.count('-v'), 3)], diff --git a/tests/gclient_utils_test.py b/tests/gclient_utils_test.py index 670775a68..2f8920327 100755 --- a/tests/gclient_utils_test.py +++ b/tests/gclient_utils_test.py @@ -13,6 +13,11 @@ import sys import time import unittest +if sys.version_info.major == 2: + from StringIO import StringIO +else: + from io import StringIO + sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from testing_support import trial_dir @@ -127,7 +132,7 @@ class CheckCallAndFilterTestCase(unittest.TestCase): 'after a short nap...') @mock.patch('subprocess2.Popen') - def testCHeckCallAndFilter_PrintStdout(self, mockPopen): + def testCheckCallAndFilter_PrintStdout(self, mockPopen): cwd = 'bleh' args = ['boo', 'foo', 'bar'] test_string = 'ahah\naccb\nallo\naddb\nāœ”' @@ -139,16 +144,76 @@ class CheckCallAndFilterTestCase(unittest.TestCase): print_stdout=True) self.assertEqual(result, test_string.encode('utf-8')) + self.assertEqual(self.stdout.getvalue().splitlines(), [ + b"________ running 'boo foo bar' in 'bleh'", + b'ahah', + b'accb', + b'allo', + b'addb', + b'\xe2\x9c\x94', + ]) + + +class AnnotatedTestCase(unittest.TestCase): + def setUp(self): + self.out = gclient_utils.MakeFileAnnotated(io.BytesIO()) + self.annotated = gclient_utils.MakeFileAnnotated( + io.BytesIO(), include_zero=True) + + def testWrite(self): + test_cases = [ + ('test string\n', b'test string\n'), + (b'test string\n', b'test string\n'), + ('āœ”\n', b'\xe2\x9c\x94\n'), + (b'\xe2\x9c\x94\n', b'\xe2\x9c\x94\n'), + ('first line\nsecondline\n', b'first line\nsecondline\n'), + (b'first line\nsecondline\n', b'first line\nsecondline\n'), + ] + + for test_input, expected_output in test_cases: + out = gclient_utils.MakeFileAnnotated(io.BytesIO()) + out.write(test_input) + self.assertEqual(out.getvalue(), expected_output) + + def testWrite_Annotated(self): + test_cases = [ + ('test string\n', b'0>test string\n'), + (b'test string\n', b'0>test string\n'), + ('āœ”\n', b'0>\xe2\x9c\x94\n'), + (b'\xe2\x9c\x94\n', b'0>\xe2\x9c\x94\n'), + ('first line\nsecondline\n', b'0>first line\n0>secondline\n'), + (b'first line\nsecondline\n', b'0>first line\n0>secondline\n'), + ] + + for test_input, expected_output in test_cases: + out = gclient_utils.MakeFileAnnotated(io.BytesIO(), include_zero=True) + out.write(test_input) + self.assertEqual(out.getvalue(), expected_output) + + def testByteByByteInput(self): + self.out.write(b'\xe2') + self.out.write(b'\x9c') + self.out.write(b'\x94') + self.out.write(b'\n') + self.out.write(b'\xe2') + self.out.write(b'\n') + self.assertEqual(self.out.getvalue(), b'\xe2\x9c\x94\n\xe2\n') + + def testByteByByteInput_Annotated(self): + self.annotated.write(b'\xe2') + self.annotated.write(b'\x9c') + self.annotated.write(b'\x94') + self.annotated.write(b'\n') + self.annotated.write(b'\xe2') + self.annotated.write(b'\n') + self.assertEqual(self.annotated.getvalue(), b'0>\xe2\x9c\x94\n0>\xe2\n') + + def testFlush_Annotated(self): + self.annotated.write(b'first line\nsecond line') + self.assertEqual(self.annotated.getvalue(), b'0>first line\n') + self.annotated.flush() self.assertEqual( - self.stdout.getvalue().splitlines(), - [ - b'________ running \'boo foo bar\' in \'bleh\'', - b'ahah', - b'accb', - b'allo', - b'addb', - b'\xe2\x9c\x94', - ]) + self.annotated.getvalue(), b'0>first line\n0>second line\n') class SplitUrlRevisionTestCase(unittest.TestCase): @@ -271,7 +336,6 @@ class GClientUtilsTest(trial_dir.TestCase): if __name__ == '__main__': - import unittest unittest.main() # vim: ts=2:sw=2:tw=80:et: