Merge pull request #658 from wjt/streaming-ReadBuffer-py3-pr-635

Test case for PR #635 (and bug #615) 

tests/test_streaming.py

@@ -120,13 +120,13 @@ class TweepyStreamTests(unittest.TestCase):
         self.assertEqual(u'Caf\xe9'.encode('utf8'), s.session.params['follow'])
 
 
-class TweepyStreamReadBuffer(unittest.TestCase):
+class TweepyStreamReadBufferTests(unittest.TestCase):
 
-    stream = """11\n{id:12345}\n\n24\n{id:23456, test:"blah"}\n"""
+    stream = six.b("""11\n{id:12345}\n\n24\n{id:23456, test:"blah"}\n""")
 
     def test_read_tweet(self):
         for length in [1, 2, 5, 10, 20, 50]:
-            buf = ReadBuffer(six.StringIO(self.stream), length)
+            buf = ReadBuffer(six.BytesIO(self.stream), length)
             self.assertEqual('11\n', buf.read_line())
             self.assertEqual('{id:12345}\n', buf.read_len(11))
             self.assertEqual('\n', buf.read_line())
@@ -157,13 +157,14 @@ class TweepyStreamReadBuffer(unittest.TestCase):
             return ""
 
         # Create a fake stream
-        stream = six.StringIO('')
+        stream = six.BytesIO(six.b(''))
 
         # Mock it's read function so it can't be called too many times
         mock_read = MagicMock(side_effect=on_read)
 
         try:
-            with patch.multiple(stream, create=True, read=mock_read, closed=True):
+            stream.close()
+            with patch.multiple(stream, create=True, read=mock_read):
                 # Now the stream can't call 'read' more than call_limit times
                 # and it looks like a requests stream that is closed
                 buf = ReadBuffer(stream, 50)
@@ -175,14 +176,14 @@ class TweepyStreamReadBuffer(unittest.TestCase):
         self.assertEqual(mock_read.call_count, 0)
 
     def test_read_unicode_tweet(self):
-        stream = '11\n{id:12345}\n\n23\n{id:23456, test:"\xe3\x81\x93"}\n\n'
+        stream = six.b('11\n{id:12345}\n\n23\n{id:23456, test:"\xe3\x81\x93"}\n\n')
         for length in [1, 2, 5, 10, 20, 50]:
-            buf = ReadBuffer(six.StringIO(stream), length)
+            buf = ReadBuffer(six.BytesIO(stream), length)
             self.assertEqual('11\n', buf.read_line())
             self.assertEqual('{id:12345}\n', buf.read_len(11))
             self.assertEqual('\n', buf.read_line())
             self.assertEqual('23\n', buf.read_line())
-            self.assertEqual('{id:23456, test:"\xe3\x81\x93"}\n', buf.read_len(23))
+            self.assertEqual(u'{id:23456, test:"\u3053"}\n', buf.read_len(23))
 
 
 class TweepyStreamBackoffTests(unittest.TestCase):

tweepy/streaming.py

@@ -7,6 +7,7 @@
 from __future__ import absolute_import, print_function
 
 import logging
+import re
 import requests
 from requests.exceptions import Timeout
 from threading import Thread
@@ -148,10 +149,11 @@ class ReadBuffer(object):
     use small chunks so it can read the length and the tweet in 2 read calls.
     """
 
-    def __init__(self, stream, chunk_size):
+    def __init__(self, stream, chunk_size, encoding='utf-8'):
         self._stream = stream
-        self._buffer = ''
+        self._buffer = six.b('')
         self._chunk_size = chunk_size
+        self._encoding = encoding
 
     def read_len(self, length):
         while not self._stream.closed:
@@ -160,7 +162,13 @@ class ReadBuffer(object):
             read_len = max(self._chunk_size, length - len(self._buffer))
             self._buffer += self._stream.read(read_len)
 
-    def read_line(self, sep='\n'):
+    def read_line(self, sep=six.b('\n')):
+        """Read the data stream until a given separator is found (default \n)
+
+        :param sep: Separator to read until. Must by of the bytes type (str in python 2,
+            bytes in python 3)
+        :return: The str of the data read until sep
+        """
         start = 0
         while not self._stream.closed:
             loc = self._buffer.find(sep, start)
@@ -173,7 +181,7 @@ class ReadBuffer(object):
     def _pop(self, length):
         r = self._buffer[:length]
         self._buffer = self._buffer[length:]
-        return r
+        return r.decode(self._encoding)
 
 
 class Stream(object):
@@ -290,7 +298,14 @@ class Stream(object):
             self.running = False
 
     def _read_loop(self, resp):
-        buf = ReadBuffer(resp.raw, self.chunk_size)
+        charset = resp.headers.get('content-type', default='')
+        enc_search = re.search('charset=(?P<enc>\S*)', charset)
+        if enc_search is not None:
+            encoding = enc_search.group('enc')
+        else:
+            encoding = 'utf-8'
+
+        buf = ReadBuffer(resp.raw, self.chunk_size, encoding=encoding)
 
         while self.running and not resp.raw.closed:
             length = 0