[tor-commits] [stem/master] CollecTor Index class

[atagar at torproject.org]

commit 860fda10124767e68e713436a4e191a2b0db4bfd
Author: Damian Johnson <atagar at torproject.org>
Date:   Sun Jul 7 11:21:00 2019 -0700

    CollecTor Index class
    
    Class to provide the methods we'll need to work with the index.
---
 stem/descriptor/collector.py      | 52 +++++++++++++++++++++++++--------------
 test/unit/descriptor/collector.py | 13 +++++-----
 2 files changed, 40 insertions(+), 25 deletions(-)

diff --git a/stem/descriptor/collector.py b/stem/descriptor/collector.py
index 2633558e..770a1296 100644
--- a/stem/descriptor/collector.py
+++ b/stem/descriptor/collector.py
@@ -121,29 +121,44 @@ def _download(url, compression, timeout, retries):
   return stem.util.str_tools._to_unicode(response)
 
 
-def _convert_index_paths(val):
+class Index(object):
+  """
+  Index of CollecTor's content.
   """
-  Key files and directories off their paths so we can work with them more
-  efficiently.
 
-  :val dict val: index to convert
+  def __init__(self, content):
+    self._str_content = content
+    self._hash_content = Index._convert_paths(json.loads(content))
 
-  :returns: index with files and directories converted to path-keyed hashes
+  def __str__(self):
+    return self._str_content
 
-  :raises: **ValueError** if unable to parse the index
-  """
+  def __iter__(self):
+    for k, v in self._hash_content.items():
+      yield k, v
+
+  @staticmethod
+  def _convert_paths(val):
+    """
+    Key files and directories off their paths so we can transverse them more
+    efficiently.
 
-  if isinstance(val, dict):
-    return dict([(k, _convert_index_paths(v)) for (k, v) in val.items()])
-  elif isinstance(val, list) and all([isinstance(entry, dict) and 'path' in entry for entry in val]):
-    contents = {}
+    :val dict val: index to convert
 
-    for entry in val:
-      contents[entry['path']] = dict((k, _convert_index_paths(v)) for (k, v) in entry.items() if k != 'path')
+    :returns: index with files and directories converted to path-keyed hashes
+    """
+
+    if isinstance(val, dict):
+      return dict([(k, Index._convert_paths(v)) for (k, v) in val.items()])
+    elif isinstance(val, list) and all([isinstance(entry, dict) and 'path' in entry for entry in val]):
+      contents = {}
 
-    return contents
-  else:
-    return val
+      for entry in val:
+        contents[entry['path']] = dict((k, Index._convert_paths(v)) for (k, v) in entry.items() if k != 'path')
+
+      return contents
+    else:
+      return val
 
 
 class CollecTor(object):
@@ -179,7 +194,8 @@ class CollecTor(object):
     """
     Provides the archives available in CollecTor.
 
-    :returns: **dict** with the archive contents
+    :returns: :class:`~stem.descriptor.collector.Index` with the archive
+      contents
 
     :raises:
       If unable to retrieve the index this provide...
@@ -192,7 +208,7 @@ class CollecTor(object):
 
     if not self._cached_index or time.time() - self._cached_index_at >= REFRESH_INDEX_RATE:
       response = _download(COLLECTOR_URL + 'index/index.json', self.compression, self.timeout, self.retries)
-      self._cached_index = _convert_index_paths(json.loads(response))
+      self._cached_index = Index(response)
       self._cached_index_at = time.time()
 
     return self._cached_index
diff --git a/test/unit/descriptor/collector.py b/test/unit/descriptor/collector.py
index 56ab19c1..2e11c89d 100644
--- a/test/unit/descriptor/collector.py
+++ b/test/unit/descriptor/collector.py
@@ -35,7 +35,7 @@ class TestCollector(unittest.TestCase):
     urlopen_mock.return_value = io.BytesIO(MINIMAL_INDEX_JSON)
 
     collector = CollecTor(compression = Compression.PLAINTEXT)
-    self.assertEqual(MINIMAL_INDEX, collector.index())
+    self.assertEqual(MINIMAL_INDEX, dict(collector.index()))
     urlopen_mock.assert_called_with('https://collector.torproject.org/index/index.json', timeout = None)
 
   @patch(URL_OPEN)
@@ -48,7 +48,7 @@ class TestCollector(unittest.TestCase):
     urlopen_mock.return_value = io.BytesIO(zlib.compress(MINIMAL_INDEX_JSON))
 
     collector = CollecTor(compression = Compression.GZIP)
-    self.assertEqual(MINIMAL_INDEX, collector.index())
+    self.assertEqual(MINIMAL_INDEX, dict(collector.index()))
     urlopen_mock.assert_called_with('https://collector.torproject.org/index/index.json.gz', timeout = None)
 
   @patch(URL_OPEN)
@@ -61,7 +61,7 @@ class TestCollector(unittest.TestCase):
     urlopen_mock.return_value = io.BytesIO(bz2.compress(MINIMAL_INDEX_JSON))
 
     collector = CollecTor(compression = Compression.BZ2)
-    self.assertEqual(MINIMAL_INDEX, collector.index())
+    self.assertEqual(MINIMAL_INDEX, dict(collector.index()))
     urlopen_mock.assert_called_with('https://collector.torproject.org/index/index.json.bz2', timeout = None)
 
   @patch(URL_OPEN)
@@ -74,7 +74,7 @@ class TestCollector(unittest.TestCase):
     urlopen_mock.return_value = io.BytesIO(lzma.compress(MINIMAL_INDEX_JSON))
 
     collector = CollecTor(compression = Compression.LZMA)
-    self.assertEqual(MINIMAL_INDEX, collector.index())
+    self.assertEqual(MINIMAL_INDEX, dict(collector.index()))
     urlopen_mock.assert_called_with('https://collector.torproject.org/index/index.json.lzma', timeout = None)
 
   @patch(URL_OPEN)
@@ -94,7 +94,7 @@ class TestCollector(unittest.TestCase):
   @patch(URL_OPEN, Mock(return_value = io.BytesIO(MINIMAL_INDEX_JSON)))
   def test_index(self):
     collector = CollecTor(compression = Compression.PLAINTEXT)
-    self.assertEqual(MINIMAL_INDEX, collector.index())
+    self.assertEqual(MINIMAL_INDEX, dict(collector.index()))
 
   @patch(URL_OPEN, Mock(return_value = io.BytesIO(b'not json')))
   def test_index_malformed_json(self):
@@ -119,6 +119,5 @@ class TestCollector(unittest.TestCase):
     with open(get_resource('collector_index.json'), 'rb') as index_file:
       urlopen_mock.return_value = io.BytesIO(index_file.read())
 
-    self.maxDiff = None
     collector = CollecTor(compression = Compression.PLAINTEXT)
-    self.assertEqual(EXAMPLE_INDEX, collector.index())
+    self.assertEqual(EXAMPLE_INDEX, dict(collector.index()))