From 26178fc3acc0ab8b86023e8a8a815bbd4b384174 Mon Sep 17 00:00:00 2001
From: Phil Ruffwind <rf@rufflewind.com>
Date: Tue, 13 Oct 2015 19:00:33 -0400
Subject: [PATCH] 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
---
 googleapiclient/discovery.py | 15 +++++++++++----
 tests/test_discovery.py      | 17 ++++++++++++++++-
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/googleapiclient/discovery.py b/googleapiclient/discovery.py
index be62cf73e13..cee56284c07 100644
--- a/googleapiclient/discovery.py
+++ b/googleapiclient/discovery.py
@@ -28,14 +28,17 @@
     '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 @@
 # 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 method(self, **kwargs):
           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 9bea84ddd14..0181bbb7e62 100644
--- a/tests/test_discovery.py
+++ b/tests/test_discovery.py
@@ -35,6 +35,7 @@
 import json
 import os
 import pickle
+import re
 import sys
 import unittest2 as unittest
 
@@ -787,7 +788,21 @@ def test_multipart_media_good_upload(self):
     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)