From e12c5b932df1c180c763a6bf1630b22928d88527 Mon Sep 17 00:00:00 2001
From: Clifton Barnes <cbarnes@smartfile.com>
Date: Tue, 6 Jun 2017 13:44:53 -0400
Subject: [PATCH] Allow the user to control the download

Allows the caller of the 'download' method to get the response object
back to download the item as needed. This allows the caller to
iteratively download in whatever chunk size is desired, to download to
whatever location is desired, etc.
---
 smartfile/__init__.py  | 16 ++++++++++++----
 test/test_smartfile.py |  3 +++
 2 files changed, 15 insertions(+), 4 deletions(-)

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))