diff --git a/smartfile/__init__.py b/smartfile/__init__.py index 44f3b3d..958e9fa 100644 --- a/smartfile/__init__.py +++ b/smartfile/__init__.py @@ -67,7 +67,10 @@ def _do_request(self, request, url, **kwargs): return response.text else: # This might be a file, so return it. - return response.raw + if kwargs.get('params', {}).get('raw', True): + return response.raw + else: + return response def _request(self, method, endpoint, id=None, **kwargs): "Handles retrying failed requests and error handling." @@ -145,14 +148,19 @@ def upload(self, filename, fileobj): arg = (filename, fileobj) return self.post('/path/data/', file=arg) - def download(self, file_to_be_downloaded): + def download(self, file_to_be_downloaded, perform_download=True): """ file_to_be_downloaded is a file-like object that has already been uploaded, you cannot download folders """ + response = self.get( + '/path/data/', file_to_be_downloaded, raw=False) + if not perform_download: + # The caller can decide how to process the download of the data + return response + # download uses shutil.copyfileobj to download, which copies # the data in chunks o = open(file_to_be_downloaded, 'wb') - return shutil.copyfileobj(self.get('/path/data/', - file_to_be_downloaded), o) + return shutil.copyfileobj(response.raw, o) def move(self, src_path, dst_path): # check destination folder for / at end diff --git a/test/test_smartfile.py b/test/test_smartfile.py index 2e19ff6..247de6c 100644 --- a/test/test_smartfile.py +++ b/test/test_smartfile.py @@ -1,6 +1,7 @@ from __future__ import absolute_import import os +import requests import unittest try: @@ -47,6 +48,8 @@ def upload(self): self.assertEquals(self.get_data()['size'], f.tell()) def download(self): + response = self.api.download(TESTFN, False) + self.assertTrue(isinstance(response, requests.Response)) self.api.download(TESTFN) self.assertEquals(self.get_data()['size'], os.path.getsize(TESTFN))