Reland "gclient_utils: buffer output as bytestrings in Annotated"

This is a reland of 5d284fdf

Always convert the input to bytes and write to sys.stdout (in Python2)
and sys.stdout.buffer (in Python 3).

Original change's description:
> gclient_utils: buffer output as bytestrings in Annotated
>
> In Python 3 byestrings and normal strings can't be concatenated.
> To fix this we buffer as bytestrings in the Annotated wrapper.
> We can't decode to a string because the output might come byte-by-byte, so it doesn't work with Unicode characters like ✔.
>
> Also had to update gclient_test.py, where double-wrapping stdout with Annotated caused made output not work and include_zero=True working caused other unintended side-effects.
>
> Example error from "fetch chromium":
> Traceback (most recent call last):
>   File "C:\Google\depot_tools\gclient_scm.py", line 1045, in _Clone
>     self._Run(clone_cmd, options, cwd=self._root_dir, retry=True,
>   File "C:\Google\depot_tools\gclient_scm.py", line 1370, in _Run
>     gclient_utils.CheckCallAndFilter(cmd, env=env, **kwargs)
>   File "C:\Google\depot_tools\gclient_utils.py", line 583, in CheckCallAndFilter
>     show_header_if_necessary(needs_header, attempt)
>   File "C:\Google\depot_tools\gclient_utils.py", line 533, in show_header_if_necessary
>     stdout_write(header.encode())
>   File "C:\Google\depot_tools\gclient_utils.py", line 391, in write
>     obj[0] += out
> TypeError: can only concatenate str (not "bytes") to str
>
> Bug: 984182
> Change-Id: If7037d30e9faf524f2405258281f6e6cd0bcdcae
> Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/1778745
> Commit-Queue: Edward Lesmes <ehmaldonado@chromium.org>
> Reviewed-by: Edward Lesmes <ehmaldonado@chromium.org>
> Auto-Submit: Raul Tambre <raul@tambre.ee>

Bug: 984182
Change-Id: Ifafb5e16a517c4c38dd4c9d5c6d4c3f994838bc9
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/1845504Reviewed-by: Anthony Polito <apolito@google.com>
Commit-Queue: Edward Lesmes <ehmaldonado@chromium.org>

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()
 
 

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)],

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: