Add _media methods and support for resumable media download.
TBR: http://codereview.appspot.com/6295077/
diff --git a/apiclient/discovery.py b/apiclient/discovery.py
index 3e1964c..d5d693e 100644
--- a/apiclient/discovery.py
+++ b/apiclient/discovery.py
@@ -52,6 +52,7 @@
from apiclient.http import MediaFileUpload
from apiclient.http import MediaUpload
from apiclient.model import JsonModel
+from apiclient.model import MediaModel
from apiclient.model import RawModel
from apiclient.schema import Schemas
from email.mime.multipart import MIMEMultipart
@@ -499,7 +500,9 @@
model = self._model
# If there is no schema for the response then presume a binary blob.
- if 'response' not in methodDesc:
+ if methodName.endswith('_media'):
+ model = MediaModel()
+ elif 'response' not in methodDesc:
model = RawModel()
headers = {}
@@ -618,8 +621,11 @@
for (name, desc) in zip(enum, enumDesc):
docs.append(' %s - %s\n' % (name, desc))
if 'response' in methodDesc:
- docs.append('\nReturns:\n An object of the form\n\n ')
- docs.append(schema.prettyPrintSchema(methodDesc['response']))
+ if methodName.endswith('_media'):
+ docs.append('\nReturns:\n The media object as a string.\n\n ')
+ else:
+ docs.append('\nReturns:\n An object of the form:\n\n ')
+ docs.append(schema.prettyPrintSchema(methodDesc['response']))
setattr(method, '__doc__', ''.join(docs))
setattr(theclass, methodName, method)
@@ -680,6 +686,10 @@
if 'methods' in resourceDesc:
for methodName, methodDesc in resourceDesc['methods'].iteritems():
createMethod(Resource, methodName, methodDesc, rootDesc)
+ # Add in _media methods. The functionality of the attached method will
+ # change when it sees that the method name ends in _media.
+ if methodDesc.get('supportsMediaDownload', False):
+ createMethod(Resource, methodName + '_media', methodDesc, rootDesc)
# Add in nested resources
if 'resources' in resourceDesc:
diff --git a/apiclient/http.py b/apiclient/http.py
index 022bdce..c3ef13b 100644
--- a/apiclient/http.py
+++ b/apiclient/http.py
@@ -76,6 +76,32 @@
return 0.0
+class MediaDownloadProgress(object):
+ """Status of a resumable download."""
+
+ def __init__(self, resumable_progress, total_size):
+ """Constructor.
+
+ Args:
+ resumable_progress: int, bytes received so far.
+ total_size: int, total bytes in complete download.
+ """
+ self.resumable_progress = resumable_progress
+ self.total_size = total_size
+
+ def progress(self):
+ """Percent of download completed, as a float.
+
+ Returns:
+ the percentage complete as a float, returning 0.0 if the total size of
+ the download is unknown.
+ """
+ if self.total_size is not None:
+ return float(self.resumable_progress) / float(self.total_size)
+ else:
+ return 0.0
+
+
class MediaUpload(object):
"""Describes a media object to upload.
@@ -268,7 +294,7 @@
return self._fd.read(length)
def to_json(self):
- """Creating a JSON representation of an instance of Credentials.
+ """Creating a JSON representation of an instance of MediaFileUpload.
Returns:
string, a JSON representation of this instance, suitable to pass to
@@ -472,6 +498,86 @@
d['_resumable'])
+class MediaIoBaseDownload(object):
+ """"Download media resources.
+
+ Note that the Python file object is compatible with io.Base and can be used
+ with this class also.
+
+
+ Example:
+ request = service.objects().get_media(
+ bucket='a_bucket_id',
+ name='smiley.png')
+
+ fh = io.FileIO('image.png', mode='wb')
+ downloader = MediaIoBaseDownload(fh, request, chunksize=1024*1024)
+
+ done = False
+ while done is False:
+ status, done = downloader.next_chunk()
+ if status:
+ print "Download %d%%." % int(status.progress() * 100)
+ print "Download Complete!"
+ """
+
+ def __init__(self, fh, request, chunksize=DEFAULT_CHUNK_SIZE):
+ """Constructor.
+
+ Args:
+ fh: io.Base or file object, The stream in which to write the downloaded
+ bytes.
+ request: apiclient.http.HttpRequest, the media request to perform in
+ chunks.
+ chunksize: int, File will be downloaded in chunks of this many bytes.
+ """
+ self.fh_ = fh
+ self.request_ = request
+ self.uri_ = request.uri
+ self.chunksize_ = chunksize
+ self.progress_ = 0
+ self.total_size_ = None
+ self.done_ = False
+
+ def next_chunk(self):
+ """Get the next chunk of the download.
+
+ Returns:
+ (status, done): (MediaDownloadStatus, boolean)
+ The value of 'done' will be True when the media has been fully
+ downloaded.
+
+ Raises:
+ apiclient.errors.HttpError if the response was not a 2xx.
+ httplib2.Error if a transport error has occured.
+ """
+ headers = {
+ 'range': 'bytes=%d-%d' % (
+ self.progress_, self.progress_ + self.chunksize_)
+ }
+ http = self.request_.http
+ http.follow_redirects = False
+
+ resp, content = http.request(self.uri_, headers=headers)
+ if resp.status in [301, 302, 303, 307, 308] and 'location' in resp:
+ self.uri_ = resp['location']
+ resp, content = http.request(self.uri_, headers=headers)
+ if resp.status in [200, 206]:
+ self.progress_ += len(content)
+ self.fh_.write(content)
+
+ if 'content-range' in resp:
+ content_range = resp['content-range']
+ length = content_range.rsplit('/', 1)[1]
+ self.total_size_ = int(length)
+
+ if self.progress_ == self.total_size_:
+ self.done_ = True
+ return MediaDownloadProgress(self.progress_, self.total_size_), self.done_
+ else:
+ raise HttpError(resp, content, self.uri_)
+
+
class HttpRequest(object):
"""Encapsulates a single HTTP request."""
@@ -1219,6 +1325,7 @@
iterable: iterable, a sequence of pairs of (headers, body)
"""
self._iterable = iterable
+ self.follow_redirects = True
def request(self, uri,
method='GET',
diff --git a/apiclient/model.py b/apiclient/model.py
index aafd88b..0e31053 100644
--- a/apiclient/model.py
+++ b/apiclient/model.py
@@ -289,6 +289,25 @@
return ''
+class MediaModel(JsonModel):
+ """Model class for requests that return Media.
+
+ Serializes and de-serializes between JSON and the Python
+ object representation of HTTP request, and returns the raw bytes
+ of the response body.
+ """
+ accept = '*/*'
+ content_type = 'application/json'
+ alt_param = 'media'
+
+ def deserialize(self, content):
+ return content
+
+ @property
+ def no_content_response(self):
+ return ''
+
+
class ProtocolBufferModel(BaseModel):
"""Model class for protocol buffers.
diff --git a/tests/data/zoo.json b/tests/data/zoo.json
index b4d7cfe..eec8b36 100644
--- a/tests/data/zoo.json
+++ b/tests/data/zoo.json
@@ -331,6 +331,7 @@
"id": "zoo.animals.get",
"httpMethod": "GET",
"description": "Get animals",
+ "supportsMediaDownload": true,
"parameters": {
"name": {
"location": "path",
diff --git a/tests/test_discovery.py b/tests/test_discovery.py
index bb77154..8d45e2d 100644
--- a/tests/test_discovery.py
+++ b/tests/test_discovery.py
@@ -623,6 +623,7 @@
self.assertEqual(expected, simplejson.loads(e.content),
'Should send an empty body when requesting the current upload status.')
+
class Next(unittest.TestCase):
def test_next_successful_none_on_no_next_page_token(self):
@@ -647,5 +648,24 @@
request = service.currentLocation().get()
+class MediaGet(unittest.TestCase):
+
+ def test_get_media(self):
+ http = HttpMock(datafile('zoo.json'), {'status': '200'})
+ zoo = build('zoo', 'v1', http)
+ request = zoo.animals().get_media(name='Lion')
+
+ parsed = urlparse.urlparse(request.uri)
+ q = parse_qs(parsed[4])
+ self.assertEqual(q['alt'], ['media'])
+ self.assertEqual(request.headers['accept'], '*/*')
+
+ http = HttpMockSequence([
+ ({'status': '200'}, 'standing in for media'),
+ ])
+ response = request.execute(http)
+ self.assertEqual('standing in for media', response)
+
+
if __name__ == '__main__':
unittest.main()
diff --git a/tests/test_http.py b/tests/test_http.py
index b8b638e..15068ba 100644
--- a/tests/test_http.py
+++ b/tests/test_http.py
@@ -27,14 +27,18 @@
import unittest
import StringIO
+from apiclient.discovery import build
from apiclient.errors import BatchError
+from apiclient.errors import HttpError
from apiclient.http import BatchHttpRequest
+from apiclient.http import HttpMock
from apiclient.http import HttpMockSequence
from apiclient.http import HttpRequest
from apiclient.http import MediaFileUpload
from apiclient.http import MediaUpload
from apiclient.http import MediaInMemoryUpload
from apiclient.http import MediaIoBaseUpload
+from apiclient.http import MediaIoBaseDownload
from apiclient.http import set_user_agent
from apiclient.model import JsonModel
from oauth2client.client import Credentials
@@ -251,6 +255,90 @@
pass
+class TestMediaIoBaseDownload(unittest.TestCase):
+
+ def setUp(self):
+ http = HttpMock(datafile('zoo.json'), {'status': '200'})
+ zoo = build('zoo', 'v1', http)
+ self.request = zoo.animals().get_media(name='Lion')
+ self.fh = StringIO.StringIO()
+
+ def test_media_io_base_download(self):
+ self.request.http = HttpMockSequence([
+ ({'status': '200',
+ 'content-range': '0-2/5'}, '123'),
+ ({'status': '200',
+ 'content-range': '3-4/5'}, '45'),
+ ])
+
+ download = MediaIoBaseDownload(
+ fh=self.fh, request=self.request, chunksize=3)
+
+ self.assertEqual(self.fh, download.fh_)
+ self.assertEqual(3, download.chunksize_)
+ self.assertEqual(0, download.progress_)
+ self.assertEqual(None, download.total_size_)
+ self.assertEqual(False, download.done_)
+ self.assertEqual(self.request.uri, download.uri_)
+
+ status, done = download.next_chunk()
+
+ self.assertEqual(self.fh.getvalue(), '123')
+ self.assertEqual(False, done)
+ self.assertEqual(3, download.progress_)
+ self.assertEqual(5, download.total_size_)
+ self.assertEqual(3, status.resumable_progress)
+
+ status, done = download.next_chunk()
+
+ self.assertEqual(self.fh.getvalue(), '12345')
+ self.assertEqual(True, done)
+ self.assertEqual(5, download.progress_)
+ self.assertEqual(5, download.total_size_)
+
+ def test_media_io_base_download_handle_redirects(self):
+ self.request.http = HttpMockSequence([
+ ({'status': '307',
+ 'location': 'https://secure.example.net/lion'}, ''),
+ ({'status': '200',
+ 'content-range': '0-2/5'}, 'abc'),
+ ])
+
+ download = MediaIoBaseDownload(
+ fh=self.fh, request=self.request, chunksize=3)
+
+ status, done = download.next_chunk()
+
+ self.assertEqual('https://secure.example.net/lion', download.uri_)
+ self.assertEqual(self.fh.getvalue(), 'abc')
+ self.assertEqual(False, done)
+ self.assertEqual(3, download.progress_)
+ self.assertEqual(5, download.total_size_)
+
+ def test_media_io_base_download_handle_4xx(self):
+ self.request.http = HttpMockSequence([
+ ({'status': '400'}, ''),
+ ])
+
+ download = MediaIoBaseDownload(
+ fh=self.fh, request=self.request, chunksize=3)
+
+ try:
+ status, done = download.next_chunk()
+ self.fail('Should raise an exception')
+ except HttpError:
+ pass
+
+ # Even after raising an exception we can pick up where we left off.
+ self.request.http = HttpMockSequence([
+ ({'status': '200',
+ 'content-range': '0-2/5'}, '123'),
+ ])
+
+ status, done = download.next_chunk()
+
+ self.assertEqual(self.fh.getvalue(), '123')
+
EXPECTED = """POST /someapi/v1/collection/?foo=bar HTTP/1.1
Content-Type: application/json
MIME-Version: 1.0
@@ -581,6 +669,5 @@
self.assertEqual({'baz': 'qux'}, callbacks.responses['2'])
-
if __name__ == '__main__':
unittest.main()