Fix non-resumable binary uploads on Python 3

 1. Generator and StringIO are replaced by BytesGenerator and BytesIO.
    If BytesGenerator doesn't exist (as is the case in Python 2), fall
    back to Generator.

 2. BytesGenerator is buggy [1] [2] and corrupts '\r' into '\n'.  To
    work around this, we implement a patched version of BytesGenerator
    that replaces ._write_lines with just .write.

The test_multipart_media_good_upload has been updated to reflect the
change.  It is also stricter now, as it matches the entire request body
against the expected form.

Note: BytesGenerator was introduced in Python 3.2.  This is OK since the
      library already demands 3.3+.

Fixes #145.

[1]: https://bugs.python.org/issue18886
[2]: https://bugs.python.org/issue19003

2 files changed, 27 insertions, 5 deletions

diff --git a/googleapiclient/discovery.py b/googleapiclient/discovery.py
index be62cf73e..cee56284c 100644
--- a/googleapiclient/discovery.py
+++ b/googleapiclient/discovery.py
@@ -28,14 +28,17 @@ __all__ = [
     'key2param',
     ]
 
-from six import StringIO
+from six import BytesIO
 from six.moves import http_client
 from six.moves.urllib.parse import urlencode, urlparse, urljoin, \
   urlunparse, parse_qsl
 
 # Standard library imports
 import copy
-from email.generator import Generator
+try:
+  from email.generator import BytesGenerator
+except ImportError:
+  from email.generator import Generator as BytesGenerator
 from email.mime.multipart import MIMEMultipart
 from email.mime.nonmultipart import MIMENonMultipart
 import json
@@ -102,6 +105,10 @@ STACK_QUERY_PARAMETER_DEFAULT_VALUE = {'type': 'string', 'location': 'query'}
 # Library-specific reserved words beyond Python keywords.
 RESERVED_WORDS = frozenset(['body'])
 
+# patch _write_lines to avoid munging '\r' into '\n'
+# ( https://bugs.python.org/issue18886 https://bugs.python.org/issue19003 )
+class _BytesGenerator(BytesGenerator):
+  _write_lines = BytesGenerator.write
 
 def fix_method_name(name):
   """Fix method names to avoid reserved word conflicts.
@@ -797,8 +804,8 @@ def createMethod(methodName, methodDesc, rootDesc, schema):
           msgRoot.attach(msg)
           # encode the body: note that we can't use `as_string`, because
           # it plays games with `From ` lines.
-          fp = StringIO()
-          g = Generator(fp, mangle_from_=False)
+          fp = BytesIO()
+          g = _BytesGenerator(fp, mangle_from_=False)
           g.flatten(msgRoot, unixfrom=False)
           body = fp.getvalue()
 
diff --git a/tests/test_discovery.py b/tests/test_discovery.py
index 9bea84ddd..0181bbb7e 100644
--- a/tests/test_discovery.py
+++ b/tests/test_discovery.py
@@ -35,6 +35,7 @@ import itertools
 import json
 import os
 import pickle
+import re
 import sys
 import unittest2 as unittest
 
@@ -787,7 +788,21 @@ class Discovery(unittest.TestCase):
     request = zoo.animals().insert(media_body=datafile('small.png'), body={})
     self.assertTrue(request.headers['content-type'].startswith(
         'multipart/related'))
-    self.assertEquals('--==', request.body[0:4])
+    with open(datafile('small.png'), 'rb') as f:
+      contents = f.read()
+    boundary = re.match(b'--=+([^=]+)', request.body).group(1)
+    self.assertEqual(
+      request.body.rstrip(b"\n"), # Python 2.6 does not add a trailing \n
+      b'--===============' + boundary + b'==\n' +
+      b'Content-Type: application/json\n' +
+      b'MIME-Version: 1.0\n\n' +
+      b'{"data": {}}\n' +
+      b'--===============' + boundary + b'==\n' +
+      b'Content-Type: image/png\n' +
+      b'MIME-Version: 1.0\n' +
+      b'Content-Transfer-Encoding: binary\n\n' +
+      contents +
+      b'\n--===============' + boundary + b'==--')
     assertUrisEqual(self,
         'https://www.googleapis.com/upload/zoo/v1/animals?uploadType=multipart&alt=json',
         request.uri)