From 44454e4c65dda5a2a811553e1cc540140c039fd5 Mon Sep 17 00:00:00 2001
From: Joe Gregorio <jcgregorio@google.com>
Date: Fri, 15 Jun 2012 08:38:53 -0400
Subject: [PATCH] Better handling of streams that report their size as 0.

Reviewed in http://codereview.appspot.com/6300093/.

Fixes issue #146.
---
 apiclient/http.py       | 15 ++++++++++++---
 tests/test_discovery.py | 26 +++++++++++++++++++++++---
 2 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/apiclient/http.py b/apiclient/http.py
index 9155cab..022bdce 100644
--- a/apiclient/http.py
+++ b/apiclient/http.py
@@ -289,7 +289,6 @@ class MediaIoBaseUpload(MediaUpload):
   Note that the Python file object is compatible with io.Base and can be used
   with this class also.
 
-
     fh = io.BytesIO('...Some data to upload...')
     media = MediaIoBaseUpload(fh, mimetype='image/png',
       chunksize=1024*1024, resumable=True)
@@ -304,7 +303,8 @@ class MediaIoBaseUpload(MediaUpload):
     """Constructor.
 
     Args:
-      fh: io.Base or file object, The source of the bytes to upload.
+      fh: io.Base or file object, The source of the bytes to upload. MUST be
+        opened in blocking mode, do not use streams opened in non-blocking mode.
       mimetype: string, Mime-type of the file. If None then a mime-type will be
         guessed from the file extension.
       chunksize: int, File will be uploaded in chunks of this many bytes. Only
@@ -320,7 +320,11 @@ class MediaIoBaseUpload(MediaUpload):
     try:
       if hasattr(self._fh, 'fileno'):
         fileno = self._fh.fileno()
-        self._size = os.fstat(fileno).st_size
+
+        # Pipes and such show up as 0 length files.
+        size = os.fstat(fileno).st_size
+        if size:
+          self._size = os.fstat(fileno).st_size
     except IOError:
       pass
 
@@ -616,6 +620,11 @@ class HttpRequest(object):
 
     data = self.resumable.getbytes(
         self.resumable_progress, self.resumable.chunksize())
+
+    # A short read implies that we are at EOF, so finish the upload.
+    if len(data) < self.resumable.chunksize():
+      size = str(self.resumable_progress + len(data))
+
     headers = {
         'Content-Range': 'bytes %d-%d/%s' % (
             self.resumable_progress, self.resumable_progress + len(data) - 1,
diff --git a/tests/test_discovery.py b/tests/test_discovery.py
index 91db712..bb77154 100644
--- a/tests/test_discovery.py
+++ b/tests/test_discovery.py
@@ -560,12 +560,32 @@ class Discovery(unittest.TestCase):
 
     # Create an upload that doesn't know the full size of the media.
     upload = MediaIoBaseUpload(
-        fh=fh, mimetype='image/png', chunksize=500, resumable=True)
+        fh=fh, mimetype='image/png', chunksize=10, resumable=True)
 
     request = zoo.animals().insert(media_body=upload, body=None)
     status, body = request.next_chunk(http)
-    self.assertEqual(body, {'Content-Range': 'bytes 0-13/*'},
-      'Should be 14 out of * bytes.')
+    self.assertEqual(body, {'Content-Range': 'bytes 0-9/*'},
+      'Should be 10 out of * bytes.')
+
+  def test_resumable_media_handle_uploads_of_unknown_size_eof(self):
+    http = HttpMockSequence([
+      ({'status': '200',
+        'location': 'http://upload.example.com'}, ''),
+      ({'status': '200'}, 'echo_request_headers_as_json'),
+      ])
+
+    self.http = HttpMock(datafile('zoo.json'), {'status': '200'})
+    zoo = build('zoo', 'v1', self.http)
+
+    fh = StringIO.StringIO('data goes here')
+
+    # Create an upload that doesn't know the full size of the media.
+    upload = MediaIoBaseUpload(
+        fh=fh, mimetype='image/png', chunksize=15, resumable=True)
+
+    request = zoo.animals().insert(media_body=upload, body=None)
+    status, body = request.next_chunk(http)
+    self.assertEqual(body, {'Content-Range': 'bytes 0-13/14'})
 
   def test_resumable_media_handle_resume_of_upload_of_unknown_size(self):
     http = HttpMockSequence([