[tor-commits] [stem/master] Download method for extrainfo descriptors

[atagar at torproject.org]

commit c109fda630132f6981d4afbba0a4e5bb878fa6ee
Author: Damian Johnson <atagar at torproject.org>
Date:   Tue Aug 6 12:18:38 2019 -0700

    Download method for extrainfo descriptors
---
 stem/descriptor/collector.py       | 39 +++++++++++++++++++++++++++++++++++++-
 test/integ/descriptor/collector.py |  8 ++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/stem/descriptor/collector.py b/stem/descriptor/collector.py
index e6e605f0..82a213f3 100644
--- a/stem/descriptor/collector.py
+++ b/stem/descriptor/collector.py
@@ -49,7 +49,8 @@ With this you can either download and read directly from CollecTor...
 ::
 
   get_instance - Provides a singleton CollecTor used for...
-    +- get_server_descriptors - published server descriptors
+    |- get_server_descriptors - published server descriptors
+    +- get_extrainfo_descriptors - published extrainfo descriptors
 
   File - Individual file residing within CollecTor
     |- read - provides descriptors from this file
@@ -57,6 +58,7 @@ With this you can either download and read directly from CollecTor...
 
   CollecTor - Downloader for descriptors from CollecTor
     |- get_server_descriptors - published server descriptors
+    |- get_extrainfo_descriptors - published extrainfo descriptors
     |
     |- index - metadata for content available from CollecTor
     +- files - files available from CollecTor
@@ -150,6 +152,17 @@ def get_server_descriptors(start = None, end = None, cache_to = None, timeout =
     yield desc
 
 
+def get_extrainfo_descriptors(start = None, end = None, cache_to = None, timeout = None, retries = 3):
+  """
+  Shorthand for
+  :func:`~stem.descriptor.collector.CollecTor.get_extrainfo_descriptors`
+  on our singleton instance.
+  """
+
+  for desc in get_instance().get_extrainfo_descriptors(start, end, cache_to, timeout, retries):
+    yield desc
+
+
 class File(object):
   """
   File within CollecTor.
@@ -393,6 +406,30 @@ class CollecTor(object):
       for desc in f.read(cache_to, timeout = timeout, retries = retries):
         yield desc
 
+  def get_extrainfo_descriptors(self, start = None, end = None, cache_to = None, timeout = None, retries = 3):
+    """
+    Provides extrainfo descriptors published during the given time range,
+    sorted oldest to newest.
+
+    :param datetime.datetime start: time range to begin with
+    :param datetime.datetime end: time range to end with
+    :param str cache_to: directory to cache archives into, if an archive is
+      available here it is not downloaded
+    :param int timeout: timeout for downloading each individual archive when
+      the connection becomes idle, no timeout applied if **None**
+    :param int retires: maximum attempts to impose on a per-archive basis
+
+    :returns: **iterator** of
+      :class:`~stem.descriptor.extrainfo_descriptor.RelayExtraInfoDescriptor`
+      for the given time range
+
+    :raises: :class:`~stem.DownloadFailed` if the download fails
+    """
+
+    for f in self.files('extra-info', start, end):
+      for desc in f.read(cache_to, timeout = timeout, retries = retries):
+        yield desc
+
   def index(self, compression = 'best'):
     """
     Provides the archives available in CollecTor.
diff --git a/test/integ/descriptor/collector.py b/test/integ/descriptor/collector.py
index 040c6b72..1a2e1e32 100644
--- a/test/integ/descriptor/collector.py
+++ b/test/integ/descriptor/collector.py
@@ -43,6 +43,14 @@ class TestCollector(unittest.TestCase):
     if not (300 < len(recent_descriptors) < 800):
       self.fail('Downloaded %i descriptors, expected 300-800' % len(recent_descriptors))  # 584 on 8/5/19
 
+  @test.require.only_run_once
+  @test.require.online
+  def test_downloading_extrainfo_descriptors(self):
+    recent_descriptors = list(stem.descriptor.collector.get_extrainfo_descriptors(start = RECENT))
+
+    if not (300 < len(recent_descriptors) < 800):
+      self.fail('Downloaded %i descriptors, expected 300-800' % len(recent_descriptors))
+
   def _test_index(self, compression):
     if compression and not compression.available:
       self.skipTest('(%s unavailable)' % compression)