[tor-commits] [stem/master] Implement RouterStatusEntry.from_str()

[atagar at torproject.org]

commit 310b5ca40b24231a0721117cfb35952b9035dcde
Author: Damian Johnson <atagar at torproject.org>
Date:   Wed Nov 21 11:36:07 2018 -0800

    Implement RouterStatusEntry.from_str()
    
    Router status entries don't have their own @type annotation, so our from_str()
    method could not provide them. Implementing their own from_str() method so
    things like RouterStatusEntryV3.from_str() will work.
---
 stem/descriptor/router_status_entry.py      | 22 ++++++++++++++++++++++
 test/unit/descriptor/router_status_entry.py | 14 ++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/stem/descriptor/router_status_entry.py b/stem/descriptor/router_status_entry.py
index 7e88c89a..b3334620 100644
--- a/stem/descriptor/router_status_entry.py
+++ b/stem/descriptor/router_status_entry.py
@@ -21,6 +21,7 @@ sources...
 
 import base64
 import binascii
+import io
 
 import stem.exit_policy
 import stem.prereq
@@ -428,6 +429,27 @@ class RouterStatusEntry(Descriptor):
     'v': _parse_v_line,
   }
 
+  @classmethod
+  def from_str(cls, content, **kwargs):
+    # Router status entries don't have their own @type annotation, so to make
+    # our subclass from_str() work we need to do the type inferencing ourself.
+
+    if cls == RouterStatusEntry:
+      raise NotImplementedError('Please use the from_str() method from RouterStatusEntry subclasses, not RouterStatusEntry itself')
+    elif 'descriptor_type' in kwargs:
+      raise ValueError("Router status entries don't have their own @type annotation. As such providing a 'descriptor_type' argument with RouterStatusEntry.from_str() does not work. Please drop the 'descriptor_type' argument when using this these subclasses' from_str() method.")
+
+    is_multiple = kwargs.pop('multiple', False)
+    validate = kwargs.pop('validate', False)
+    results = list(_parse_file(io.BytesIO(stem.util.str_tools._to_bytes(content)), validate, cls, **kwargs))
+
+    if is_multiple:
+      return results
+    elif len(results) == 1:
+      return results[0]
+    else:
+      raise ValueError("Descriptor.from_str() expected a single descriptor, but had %i instead. Please include 'multiple = True' if you want a list of results instead." % len(results))
+
   def __init__(self, content, validate = False, document = None):
     """
     Parse a router descriptor in a network status document.
diff --git a/test/unit/descriptor/router_status_entry.py b/test/unit/descriptor/router_status_entry.py
index 05b4dee6..c428d6b2 100644
--- a/test/unit/descriptor/router_status_entry.py
+++ b/test/unit/descriptor/router_status_entry.py
@@ -20,6 +20,7 @@ from test.unit.descriptor import (
 )
 
 from stem.descriptor.router_status_entry import (
+  RouterStatusEntry,
   RouterStatusEntryV2,
   RouterStatusEntryV3,
   RouterStatusEntryMicroV3,
@@ -83,6 +84,19 @@ def vote_document():
 
 
 class TestRouterStatusEntry(unittest.TestCase):
+  def test_from_str(self):
+    """
+    Exercise our RouterStatusEntry.from_str().
+    """
+
+    desc = RouterStatusEntryV3.create()
+    content = desc.get_bytes()
+
+    self.assertEqual(desc, RouterStatusEntryV3.from_str(content))
+
+    self.assertRaisesWith(NotImplementedError, 'Please use the from_str() method from RouterStatusEntry subclasses, not RouterStatusEntry itself', RouterStatusEntry.from_str, content)
+    self.assertRaisesWith(ValueError, "Router status entries don't have their own @type annotation. As such providing a 'descriptor_type' argument with RouterStatusEntry.from_str() does not work. Please drop the 'descriptor_type' argument when using this these subclasses' from_str() method.", RouterStatusEntryV3.from_str, content, descriptor_type = 'network-status-consensus-3 1.0')
+
   def test_fingerprint_decoding(self):
     """
     Tests for the _base64_to_hex() helper.