#33 Support for streaming requests

3 files changed, 29 insertions, 1 deletions

diff --git a/docs/user_guide.rst b/docs/user_guide.rst
index f089ba6..37580a8 100644
--- a/docs/user_guide.rst
+++ b/docs/user_guide.rst
@@ -112,6 +112,7 @@ It can be used, for example, for request throttling with help of ``requests`` ho
 
 .. seealso:: `example.py <https://github.com/reclosedev/requests-cache/blob/master/example.py>`_
 
+.. note:: requests_cache prefetchs response content, be aware if your code uses streaming requests.
 
 .. _persistence:
 
diff --git a/requests_cache/backends/base.py b/requests_cache/backends/base.py
index 3645754..3d02527 100644
--- a/requests_cache/backends/base.py
+++ b/requests_cache/backends/base.py
@@ -10,6 +10,7 @@
 from datetime import datetime
 import hashlib
 from copy import copy
+from io import BytesIO
 
 import requests
 
@@ -140,7 +141,7 @@ class BaseCache(object):
             value = copy(value)
             value.hooks = []
         elif name == 'raw':
-            result = _Store()
+            result = _RawStore()
             for field in self._raw_response_attrs:
                 setattr(result, field, getattr(value, field, None))
             value = result
@@ -158,6 +159,7 @@ class BaseCache(object):
         result = requests.Response()
         for field in self._response_attrs:
             setattr(result, field, getattr(response, field, None))
+        result.raw._cached_content_ = result.content
         seen[id(response)] = result
         result.history = tuple(self.restore_response(r, seen) for r in response.history)
         return result
@@ -176,9 +178,20 @@ class BaseCache(object):
 
 # used for saving response attributes
 class _Store(object):
+    pass
+
+
+class _RawStore(object):
+    # noop for cached response
     def release_conn(self):
         pass
 
+    # for streaming requests support
+    def read(self, chunk_size=1):
+        if not hasattr(self, "_io_with_content_"):
+            self._io_with_content_ = BytesIO(self._cached_content_)
+        return self._io_with_content_.read(chunk_size)
+
 
 def _to_bytes(s, encoding='utf-8'):
     if is_py2 or isinstance(s, bytes):
diff --git a/tests/test_cache.py b/tests/test_cache.py
index aa7d462..f4728d7 100644
--- a/tests/test_cache.py
+++ b/tests/test_cache.py
@@ -34,6 +34,7 @@ def httpbin(*suffix):
     """Returns url for HTTPBIN resource."""
     return HTTPBIN_URL + '/'.join(suffix)
 
+
 class CacheTestCase(unittest.TestCase):
 
     def setUp(self):
@@ -286,6 +287,19 @@ class CacheTestCase(unittest.TestCase):
         self.assertTrue(s.post(url, data=sorted(params.items())).from_cache)
         self.assertFalse(s.post(url, data=sorted(params.items(), reverse=True)).from_cache)
 
+    def test_stream_requests_support(self):
+        n = 100
+        url = httpbin("stream/%s" % n)
+        r = self.s.get(url, stream=True)
+        lines = list(r.iter_lines())
+        self.assertEquals(len(lines), n)
+
+        for i in range(2):
+            r = self.s.get(url, stream=True)
+            self.assertTrue(r.from_cache)
+            cached_lines = list(r.iter_lines())
+            self.assertEquals(cached_lines, lines)
+
 
 if __name__ == '__main__':
     unittest.main()