commit cf3c55ef6aec640711e26be4a4f985850612c29c Author: Damian Johnson atagar@torproject.org Date: Mon Sep 18 10:21:41 2017 -0700
Remove bytes normalization from recv_message
Sockets always provide bytes, but we still normalized the type because tests and a few methods like ControlMessage.from_str() provided it a StringIO instead. This is silly. We can provide BytesIOs instead. Our recv_message is in the critical path so it's important to keep it as minimal as possible.
This might break our tests for python 2.6. Not sure - don't have that python verison and getting it on modern systems is a pita. I'll check later. 2.6 has the io module, but unclear if it has BytesIO. If so we might need to tweak our imports later. --- stem/control.py | 4 ++-- stem/response/__init__.py | 2 +- stem/socket.py | 6 +----- test/unit/response/control_message.py | 15 ++++++--------- 4 files changed, 10 insertions(+), 17 deletions(-)
diff --git a/stem/control.py b/stem/control.py index 2f1268ef..fd5f9f87 100644 --- a/stem/control.py +++ b/stem/control.py @@ -3298,7 +3298,7 @@ class Controller(BaseController): response = self.get_info('circuit-status')
for circ in response.splitlines(): - circ_message = stem.socket.recv_message(io.StringIO(stem.util.str_tools._to_unicode('650 CIRC ' + circ + '\r\n'))) + circ_message = stem.socket.recv_message(io.BytesIO(stem.util.str_tools._to_unicode('650 CIRC %s\r\n' % circ))) stem.response.convert('EVENT', circ_message, arrived_at = 0) circuits.append(circ_message)
@@ -3481,7 +3481,7 @@ class Controller(BaseController): response = self.get_info('stream-status')
for stream in response.splitlines(): - message = stem.socket.recv_message(io.StringIO(stem.util.str_tools._to_unicode('650 STREAM ' + stream + '\r\n'))) + message = stem.socket.recv_message(io.BytesIO(stem.util.str_tools._to_bytes('650 STREAM %s\r\n' % stream))) stem.response.convert('EVENT', message, arrived_at = 0) streams.append(message)
diff --git a/stem/response/__init__.py b/stem/response/__init__.py index eadb25c7..e0de0f5c 100644 --- a/stem/response/__init__.py +++ b/stem/response/__init__.py @@ -156,7 +156,7 @@ class ControlMessage(object):
content = re.sub('([\r]?)\n', '\r\n', content)
- msg = stem.socket.recv_message(io.StringIO(stem.util.str_tools._to_unicode(content))) + msg = stem.socket.recv_message(io.BytesIO(stem.util.str_tools._to_bytes(content)))
if msg_type is not None: convert(msg_type, msg, **kwargs) diff --git a/stem/socket.py b/stem/socket.py index ca56ea69..9817c129 100644 --- a/stem/socket.py +++ b/stem/socket.py @@ -530,11 +530,7 @@ def recv_message(control_file):
while True: try: - # From a real socket readline() would always provide bytes, but during - # tests we might be given a StringIO in which case it's unicode under - # python 3.x. - - line = stem.util.str_tools._to_bytes(control_file.readline()) + line = control_file.readline() except AttributeError: # if the control_file has been closed then we will receive: # AttributeError: 'NoneType' object has no attribute 'recv' diff --git a/test/unit/response/control_message.py b/test/unit/response/control_message.py index c0fb2f32..8b05d6a8 100644 --- a/test/unit/response/control_message.py +++ b/test/unit/response/control_message.py @@ -2,17 +2,14 @@ Unit tests for the stem.response.ControlMessage parsing and class. """
+import io import socket import unittest
-try: - from StringIO import StringIO -except: - from io import StringIO - import stem.socket import stem.response import stem.response.getinfo +import stem.util.str_tools
OK_REPLY = '250 OK\r\n'
@@ -128,7 +125,7 @@ class TestControlMessage(unittest.TestCase): for index, line in enumerate(infonames_lines): # replace the CRLF for the line infonames_lines[index] = line.rstrip('\r\n') + '\n' - test_socket_file = StringIO(''.join(infonames_lines)) + test_socket_file = io.BytesIO(stem.util.str_tools._to_bytes(''.join(infonames_lines))) self.assertRaises(stem.ProtocolError, stem.socket.recv_message, test_socket_file)
# puts the CRLF back @@ -154,8 +151,8 @@ class TestControlMessage(unittest.TestCase): # - this is part of the message prefix # - this is disrupting the line ending
- self.assertRaises(stem.ProtocolError, stem.socket.recv_message, StringIO(removal_test_input)) - self.assertRaises(stem.ProtocolError, stem.socket.recv_message, StringIO(replacement_test_input)) + self.assertRaises(stem.ProtocolError, stem.socket.recv_message, io.BytesIO(stem.util.str_tools._to_bytes(removal_test_input))) + self.assertRaises(stem.ProtocolError, stem.socket.recv_message, io.BytesIO(stem.util.str_tools._to_bytes(replacement_test_input))) else: # otherwise the data will be malformed, but this goes undetected self._assert_message_parses(removal_test_input) @@ -179,7 +176,7 @@ class TestControlMessage(unittest.TestCase): stem.response.ControlMessage for the given input """
- message = stem.socket.recv_message(StringIO(controller_reply)) + message = stem.socket.recv_message(io.BytesIO(stem.util.str_tools._to_bytes(controller_reply)))
# checks that the raw_content equals the input value self.assertEqual(controller_reply, message.raw_content())