[tor-commits] [stem/master] Support for getting byte content from get_info

[atagar at torproject.org]

commit a94f1f0a51faacf1fda5350cdfd489d45b41adc4
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon May 20 09:35:45 2013 -0700

    Support for getting byte content from get_info
    
    Adding an optional flag to get byte content from get_info rather than a str.
    This only effects python 3.x.
---
 stem/control.py               |    9 ++++++++-
 stem/response/getinfo.py      |   15 +++++++++------
 test/unit/response/getinfo.py |   13 +++++++------
 3 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/stem/control.py b/stem/control.py
index a4dde55..9ea77f8 100644
--- a/stem/control.py
+++ b/stem/control.py
@@ -711,7 +711,7 @@ class Controller(BaseController):
     import stem.connection
     stem.connection.authenticate(self, *args, **kwargs)
 
-  def get_info(self, params, default = UNDEFINED):
+  def get_info(self, params, default = UNDEFINED, get_bytes = False):
     """
     Queries the control socket for the given GETINFO option. If provided a
     default then that's returned if the GETINFO option is undefined or the
@@ -720,6 +720,7 @@ class Controller(BaseController):
 
     :param str,list params: GETINFO option or options to be queried
     :param object default: response if the query fails
+    :param bool get_bytes: provides **bytes** values rather than a **str** under python 3.x
 
     :returns:
       Response depends upon how we were called as follows...
@@ -781,6 +782,12 @@ class Controller(BaseController):
       response = self.msg("GETINFO %s" % " ".join(params))
       stem.response.convert("GETINFO", response)
       response.assert_matches(params)
+
+      # usually we want unicode values under python 3.x
+
+      if stem.prereq.is_python_3() and not get_bytes:
+        response.entries = dict((k, stem.util.str_tools._to_unicode(v)) for (k, v) in response.entries.items())
+
       reply.update(response.entries)
 
       if self.is_caching_enabled():
diff --git a/stem/response/getinfo.py b/stem/response/getinfo.py
index 1d749e4..998041d 100644
--- a/stem/response/getinfo.py
+++ b/stem/response/getinfo.py
@@ -9,7 +9,7 @@ class GetInfoResponse(stem.response.ControlMessage):
   """
   Reply for a GETINFO query.
 
-  :var dict entries: mapping between the queried options and their values
+  :var dict entries: mapping between the queried options and their bytes values
   """
 
   def _parse_message(self):
@@ -26,9 +26,9 @@ class GetInfoResponse(stem.response.ControlMessage):
     # 250 OK
 
     self.entries = {}
-    remaining_lines = list(self)
+    remaining_lines = [content for (code, div, content) in self.content(get_bytes = True)]
 
-    if not self.is_ok() or not remaining_lines.pop() == "OK":
+    if not self.is_ok() or not remaining_lines.pop() == b"OK":
       unrecognized_keywords = []
       for code, _, line in self.content():
         if code == '552' and line.startswith("Unrecognized key \"") and line.endswith("\""):
@@ -41,15 +41,18 @@ class GetInfoResponse(stem.response.ControlMessage):
 
     while remaining_lines:
       try:
-        key, value = remaining_lines.pop(0).split("=", 1)
+        key, value = remaining_lines.pop(0).split(b"=", 1)
       except ValueError:
         raise stem.ProtocolError("GETINFO replies should only contain parameter=value mappings:\n%s" % self)
 
+      if stem.prereq.is_python_3():
+        key = stem.util.str_tools._to_unicode(key)
+
       # if the value is a multiline value then it *must* be of the form
       # '<key>=\n<value>'
 
-      if "\n" in value:
-        if not value.startswith("\n"):
+      if b"\n" in value:
+        if not value.startswith(b"\n"):
           raise stem.ProtocolError("GETINFO response contained a multi-line value that didn't start with a newline:\n%s" % self)
 
         value = value[1:]
diff --git a/test/unit/response/getinfo.py b/test/unit/response/getinfo.py
index b6d2c8d..082bafc 100644
--- a/test/unit/response/getinfo.py
+++ b/test/unit/response/getinfo.py
@@ -7,6 +7,7 @@ import unittest
 import stem.response
 import stem.response.getinfo
 import stem.socket
+import stem.util.str_tools
 
 from test import mocking
 
@@ -72,7 +73,7 @@ class TestGetInfoResponse(unittest.TestCase):
 
     control_message = mocking.get_message(SINGLE_RESPONSE)
     stem.response.convert("GETINFO", control_message)
-    self.assertEqual({"version": "0.2.3.11-alpha-dev"}, control_message.entries)
+    self.assertEqual({"version": b"0.2.3.11-alpha-dev"}, control_message.entries)
 
   def test_batch_response(self):
     """
@@ -83,9 +84,9 @@ class TestGetInfoResponse(unittest.TestCase):
     stem.response.convert("GETINFO", control_message)
 
     expected = {
-      "version": "0.2.3.11-alpha-dev",
-      "address": "67.137.76.214",
-      "fingerprint": "5FDE0422045DF0E1879A3738D09099EB4A0C5BA0",
+      "version": b"0.2.3.11-alpha-dev",
+      "address": b"67.137.76.214",
+      "fingerprint": b"5FDE0422045DF0E1879A3738D09099EB4A0C5BA0",
     }
 
     self.assertEqual(expected, control_message.entries)
@@ -100,8 +101,8 @@ class TestGetInfoResponse(unittest.TestCase):
     stem.response.convert("GETINFO", control_message)
 
     expected = {
-      "version": "0.2.3.11-alpha-dev (git-ef0bc7f8f26a917c)",
-      "config-text": "\n".join(MULTILINE_RESPONSE.splitlines()[2:8]),
+      "version": b"0.2.3.11-alpha-dev (git-ef0bc7f8f26a917c)",
+      "config-text": b"\n".join(stem.util.str_tools._to_bytes(MULTILINE_RESPONSE).splitlines()[2:8]),
     }
 
     self.assertEqual(expected, control_message.entries)