commit 860fda10124767e68e713436a4e191a2b0db4bfd Author: Damian Johnson <atagar@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()))