Make ResponseReader more streamlike, so that it can be wrapped in a BuffererReader

rm-hull · rm-hull · commit c0ac736e8867 · 2014-03-13T13:02:39.000Z
diff --git a/splunklib/binding.py b/splunklib/binding.py
@@ -29,6 +29,7 @@
 import socket
 import ssl
 import urllib
+import io
 
 from datetime import datetime
 from functools import wraps
@@ -1110,7 +1111,7 @@ def request(self, url, message, **kwargs):
 
 
 # Converts an httplib response into a file-like object.
-class ResponseReader(object):
+class ResponseReader(io.RawIOBase):
     """This class provides a file-like interface for :class:`httplib` responses.
 
     The ``ResponseReader`` class is intended to be a layer to unify the different
@@ -1164,6 +1165,23 @@ def read(self, size = None):
         r = r + self._response.read(size)
         return r
 
+    def readable(self):
+        """ Indicates that the response reader is readable."""
+        return True
+
+    def readinto(self, byte_array):
+        """ Read data into a byte array, upto the size of the byte array.
+
+        :param byte_array: A byte array/memory view to pour bytes into.
+        :type byte_array: ``bytearray`` or ``memoryview``
+
+        """
+        max_size = len(byte_array)
+        data = self.read(max_size)
+        bytes_read = len(data)
+        byte_array[:bytes_read] = data
+        return bytes_read
+
 
 def handler(key_file=None, cert_file=None, timeout=None):
     """This class returns an instance of the default HTTP request handler using
diff --git a/splunklib/results.py b/splunklib/results.py
@@ -12,15 +12,15 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
-"""The **splunklib.results** module provides a streaming XML reader for Splunk 
+"""The **splunklib.results** module provides a streaming XML reader for Splunk
 search results.
 
-Splunk search results can be returned in a variety of formats including XML, 
-JSON, and CSV. To make it easier to stream search results in XML format, they 
-are returned as a stream of XML *fragments*, not as a single XML document. This 
-module supports incrementally reading one result record at a time from such a 
+Splunk search results can be returned in a variety of formats including XML,
+JSON, and CSV. To make it easier to stream search results in XML format, they
+are returned as a stream of XML *fragments*, not as a single XML document. This
+module supports incrementally reading one result record at a time from such a
 result stream. This module also provides a friendly iterator-based interface for
-accessing search results while avoiding buffering the result set, which can be 
+accessing search results while avoiding buffering the result set, which can be
 very large.
 
 To use the reader, instantiate :class:`ResultsReader` on a search result stream
@@ -65,13 +65,13 @@ class Message(object):
     def __init__(self, type_, message):
         self.type = type_
         self.message = message
-    
+
     def __repr__(self):
         return "%s: %s" % (self.type, self.message)
-    
+
     def __eq__(self, other):
         return (self.type, self.message) == (other.type, other.message)
-    
+
     def __hash__(self):
         return hash((self.type, self.message))
 
@@ -149,18 +149,18 @@ def read(self, n=None):
         return response
 
 class ResultsReader(object):
-    """This class returns dictionaries and Splunk messages from an XML results 
+    """This class returns dictionaries and Splunk messages from an XML results
     stream.
 
-    ``ResultsReader`` is iterable, and returns a ``dict`` for results, or a 
-    :class:`Message` object for Splunk messages. This class has one field, 
-    ``is_preview``, which is ``True`` when the results are a preview from a 
+    ``ResultsReader`` is iterable, and returns a ``dict`` for results, or a
+    :class:`Message` object for Splunk messages. This class has one field,
+    ``is_preview``, which is ``True`` when the results are a preview from a
     running search, or ``False`` when the results are from a completed search.
 
-    This function has no network activity other than what is implicit in the 
+    This function has no network activity other than what is implicit in the
     stream it operates on.
 
-    :param `stream`: The stream to read from (any object that supports 
+    :param `stream`: The stream to read from (any object that supports
         ``.read()``).
 
     **Example**::
@@ -224,7 +224,7 @@ def _parse_results(self, stream):
                         yield result
                         result = None
                         elem.clear()
-    
+
                 elif elem.tag == 'field' and result is not None:
                     # We need the 'result is not None' check because
                     # 'field' is also the element name in the <meta>
@@ -244,11 +244,11 @@ def _parse_results(self, stream):
                         # arbitrarily large memory intead of
                         # streaming.
                         elem.clear()
-    
+
                 elif elem.tag in ('text', 'v') and event == 'end':
                     values.append(elem.text.encode('utf8'))
                     elem.clear()
-    
+
                 elif elem.tag == 'msg':
                     if event == 'start':
                         msg_type = elem.attrib['type']
diff --git a/tests/test_binding.py b/tests/test_binding.py
@@ -66,6 +66,10 @@ def test_empty(self):
         self.assertTrue(response.empty)
         self.assertEqual(response.peek(10), "")
         self.assertEqual(response.read(10), "")
+
+        arr = bytearray(10)
+        self.assertEqual(response.readinto(arr), 0)
+        self.assertEqual(arr, bytearray(10))
         self.assertTrue(response.empty)
 
     def test_read_past_end(self):
@@ -87,6 +91,42 @@ def test_read_partial(self):
         self.assertTrue(response.empty)
         self.assertEqual(response.read(), '')
 
+    def test_readable(self):
+        txt = "abcd"
+        response = binding.ResponseReader(StringIO(txt))
+        self.assertTrue(response.readable())
+
+    def test_readinto_bytearray(self):
+        txt = "Checking readinto works as expected"
+        response = binding.ResponseReader(StringIO(txt))
+        arr = bytearray(10)
+        self.assertEqual(response.readinto(arr), 10)
+        self.assertEqual(arr[:10], "Checking r")
+        self.assertEqual(response.readinto(arr), 10)
+        self.assertEqual(arr[:10], "eadinto wo")
+        self.assertEqual(response.readinto(arr), 10)
+        self.assertEqual(arr[:10], "rks as exp")
+        self.assertEqual(response.readinto(arr), 5)
+        self.assertEqual(arr[:5], "ected")
+        self.assertTrue(response.empty)
+
+    def test_readinto_memoryview(self):
+        txt = "Checking readinto works as expected"
+        response = binding.ResponseReader(StringIO(txt))
+        arr = bytearray(10)
+        mv = memoryview(arr)
+        self.assertEqual(response.readinto(mv), 10)
+        self.assertEqual(arr[:10], "Checking r")
+        self.assertEqual(response.readinto(mv), 10)
+        self.assertEqual(arr[:10], "eadinto wo")
+        self.assertEqual(response.readinto(mv), 10)
+        self.assertEqual(arr[:10], "rks as exp")
+        self.assertEqual(response.readinto(mv), 5)
+        self.assertEqual(arr[:5], "ected")
+        self.assertTrue(response.empty)
+
+
+
 class TestUrlEncoded(BindingTestCase):
     def test_idempotent(self):
         a = UrlEncoded('abc')
@@ -95,7 +135,7 @@ def test_idempotent(self):
     def test_append(self):
         self.assertEqual(UrlEncoded('a') + UrlEncoded('b'),
                          UrlEncoded('ab'))
-    
+
     def test_append_string(self):
         self.assertEqual(UrlEncoded('a') + '%',
                          UrlEncoded('a%'))
@@ -111,15 +151,15 @@ def test_chars(self):
         for char, code in [(' ', '%20'),
                            ('"', '%22'),
                            ('%', '%25')]:
-            self.assertEqual(UrlEncoded(char), 
+            self.assertEqual(UrlEncoded(char),
                              UrlEncoded(code, skip_encode=True))
 
     def test_repr(self):
         self.assertEqual(repr(UrlEncoded('% %')), "UrlEncoded('% %')")
 
 class TestAuthority(unittest.TestCase):
     def test_authority_default(self):
-        self.assertEqual(binding._authority(), 
+        self.assertEqual(binding._authority(),
                          "https://localhost:8089")
 
     def test_ipv4_host(self):
@@ -137,7 +177,7 @@ def test_ipv6_host(self):
     def test_all_fields(self):
         self.assertEqual(
             binding._authority(
-                scheme="http", 
+                scheme="http",
                 host="splunk.utopia.net",
                 port="471"),
             "http://splunk.utopia.net:471")
@@ -172,7 +212,7 @@ def test_user_without_role_fails(self):
 
     def test_create_user(self):
         response = self.context.post(
-            PATH_USERS, name=self.username, 
+            PATH_USERS, name=self.username,
             password=self.password, roles=self.roles)
         self.assertEqual(response.status, 201)
 
@@ -265,7 +305,7 @@ def test_without_autologin(self):
         self.context.autologin = False
         self.assertEqual(self.context.get("/services").status, 200)
         self.context.logout()
-        self.assertRaises(AuthenticationError, 
+        self.assertRaises(AuthenticationError,
                           self.context.get, "/services")
 
 class TestAbspath(BindingTestCase):
@@ -403,7 +443,7 @@ def isatom(body):
 class TestPluggableHTTP(testlib.SDKTestCase):
     # Verify pluggable HTTP reqeust handlers.
     def test_handlers(self):
-        paths = ["/services", "authentication/users", 
+        paths = ["/services", "authentication/users",
                  "search/jobs"]
         handlers = [binding.handler(),  # default handler
                     urllib2_handler]
@@ -421,7 +461,7 @@ def test_logout(self):
         response = self.context.get("/services")
         self.assertEqual(response.status, 200)
         self.context.logout()
-        self.assertRaises(AuthenticationError, 
+        self.assertRaises(AuthenticationError,
                           self.context.get, "/services")
         self.assertRaises(AuthenticationError,
                           self.context.post, "/services")
@@ -512,11 +552,11 @@ def test_preexisting_token(self):
         opts["token"] = token
         opts["username"] = "boris the mad baboon"
         opts["password"] = "nothing real"
-        
+
         newContext = binding.Context(**opts)
         response = newContext.get("/services")
         self.assertEqual(response.status, 200)
-        
+
         socket = newContext.connect()
         socket.write("POST %s HTTP/1.1\r\n" % \
                          self.context._abspath("some/path/to/post/to"))
diff --git a/tests/test_results.py b/tests/test_results.py
@@ -18,14 +18,15 @@
 import testlib
 from time import sleep
 import splunklib.results as results
+import io
 
 
 class ResultsTestCase(testlib.SDKTestCase):
     def test_read_from_empty_result_set(self):
         job = self.service.jobs.create("search index=_internal_does_not_exist | head 2")
         while not job.is_done():
             sleep(0.5)
-        self.assertEquals(0, len(list(results.ResultsReader(job.results()))))
+        self.assertEquals(0, len(list(results.ResultsReader(io.BufferedReader(job.results())))))
 
     def test_read_normal_results(self):
         xml_text = """
@@ -107,9 +108,9 @@ def test_read_normal_results(self):
                 'sum(kb)': '5838.935649',
             },
         ]
-        
+
         self.assert_parsed_results_equals(xml_text, expected_results)
-    
+
     def test_read_raw_field(self):
         xml_text = """
 <?xml version='1.0' encoding='UTF-8'?>
@@ -129,9 +130,9 @@ def test_read_raw_field(self):
                 '_raw': '07-13-2012 09:27:27.307 -0700 INFO  Metrics - group=search_concurrency, system total, active_hist_searches=0, active_realtime_searches=0',
             },
         ]
-        
+
         self.assert_parsed_results_equals(xml_text, expected_results)
-    
+
     def assert_parsed_results_equals(self, xml_text, expected_results):
         results_reader = results.ResultsReader(StringIO(xml_text))
         actual_results = [x for x in results_reader]
