[tor-commits] [stem/master] Move zstd and lzma checks to prereq module

atagar at torproject.org atagar at torproject.org
Tue May 8 20:20:09 UTC 2018


commit bb260c0ab9bd7d8fad08657b21a36a186e7a5fe0
Author: Damian Johnson <atagar at torproject.org>
Date:   Sat May 5 13:41:36 2018 -0700

    Move zstd and lzma checks to prereq module
    
    This is the spot where we do these prereq checks. Yay, remote module's finally
    under the 1k line threshold.
---
 stem/descriptor/remote.py      | 38 +++++--------------------------
 stem/prereq.py                 | 51 ++++++++++++++++++++++++++++++++++++++++--
 test/settings.cfg              |  1 +
 test/unit/descriptor/remote.py | 12 +++++-----
 4 files changed, 62 insertions(+), 40 deletions(-)

diff --git a/stem/descriptor/remote.py b/stem/descriptor/remote.py
index 301e89d9..1b6a7ea6 100644
--- a/stem/descriptor/remote.py
+++ b/stem/descriptor/remote.py
@@ -113,29 +113,6 @@ try:
 except ImportError:
   import urllib2 as urllib
 
-try:
-  # added in python 3.3
-  import lzma
-  LZMA_SUPPORTED = True
-except ImportError:
-  LZMA_SUPPORTED = False
-
-try:
-  # We use the suggested python zstd library...
-  #
-  #   https://pypi.python.org/pypi/zstandard
-  #
-  # Unfortunately this installs as a zstd module which can be confused with...
-  #
-  #   https://pypi.python.org/pypi/zstd
-  #
-  # As such checking for the specific decompression class we'll need.
-
-  import zstd
-  ZSTD_SUPPORTED = hasattr(zstd, 'ZstdDecompressor')
-except ImportError:
-  ZSTD_SUPPORTED = False
-
 Compression = stem.util.enum.Enum(
   ('PLAINTEXT', 'identity'),
   ('GZIP', 'gzip'),  # can also be 'deflate'
@@ -143,9 +120,6 @@ Compression = stem.util.enum.Enum(
   ('LZMA', 'x-tor-lzma'),
 )
 
-ZSTD_UNAVAILABLE_MSG = 'ZSTD compression requires the zstandard module (https://pypi.python.org/pypi/zstandard)'
-LZMA_UNAVAILABLE_MSG = 'LZMA compression requires the lzma module (https://docs.python.org/3/library/lzma.html)'
-
 # Tor has a limited number of descriptors we can fetch explicitly by their
 # fingerprint or hashes due to a limit on the url length by squid proxies.
 
@@ -340,9 +314,10 @@ def _decompress(data, encoding):
   elif encoding in (Compression.GZIP, 'deflate'):
     return zlib.decompress(data, zlib.MAX_WBITS | 32).strip()
   elif encoding == Compression.ZSTD:
-    if not ZSTD_SUPPORTED:
+    if not stem.prereq.is_zstd_available():
       raise ImportError('Decompressing zstd data requires https://pypi.python.org/pypi/zstandard')
 
+    import zstd
     output_buffer = io.BytesIO()
 
     with zstd.ZstdDecompressor().write_to(output_buffer) as decompressor:
@@ -350,9 +325,10 @@ def _decompress(data, encoding):
 
     return output_buffer.getvalue().strip()
   elif encoding == Compression.LZMA:
-    if not LZMA_SUPPORTED:
+    if not stem.prereq.is_lzma_available():
       raise ImportError('Decompressing lzma data requires https://docs.python.org/3/library/lzma.html')
 
+    import lzma
     return lzma.decompress(data).strip()
   else:
     raise ValueError("'%s' isn't a recognized type of encoding" % encoding)
@@ -528,12 +504,10 @@ class Query(object):
       if isinstance(compression, str):
         compression = [compression]  # caller provided only a single option
 
-      if Compression.ZSTD in compression and not ZSTD_SUPPORTED:
-        log.log_once('stem.descriptor.remote.zstd_unavailable', log.INFO, ZSTD_UNAVAILABLE_MSG)
+      if Compression.ZSTD in compression and not stem.prereq.is_zstd_available():
         compression.remove(Compression.ZSTD)
 
-      if Compression.LZMA in compression and not LZMA_SUPPORTED:
-        log.log_once('stem.descriptor.remote.lzma_unavailable', log.INFO, LZMA_UNAVAILABLE_MSG)
+      if Compression.LZMA in compression and not stem.prereq.is_lzma_available():
         compression.remove(Compression.LZMA)
 
       if not compression:
diff --git a/stem/prereq.py b/stem/prereq.py
index 16437c23..6e230007 100644
--- a/stem/prereq.py
+++ b/stem/prereq.py
@@ -16,6 +16,8 @@ Checks for stem dependencies. We require python 2.6 or greater (including the
   is_python_3 - checks if python 3.0 or later is available
   is_sqlite_available - checks if the sqlite3 module is available
   is_crypto_available - checks if the cryptography module is available
+  is_zstd_available - checks if the zstd module is available
+  is_lzma_available - checks if the lzma module is available
   is_mock_available - checks if the mock module is available
 """
 
@@ -29,6 +31,8 @@ except ImportError:
   from stem.util.lru_cache import lru_cache
 
 CRYPTO_UNAVAILABLE = "Unable to import the cryptography module. Because of this we'll be unable to verify descriptor signature integrity. You can get cryptography from: https://pypi.python.org/pypi/cryptography"
+ZSTD_UNAVAILABLE = 'ZSTD compression requires the zstandard module (https://pypi.python.org/pypi/zstandard)'
+LZMA_UNAVAILABLE = 'LZMA compression requires the lzma module (https://docs.python.org/3/library/lzma.html)'
 PYNACL_UNAVAILABLE = "Unable to import the pynacl module. Because of this we'll be unable to verify descriptor ed25519 certificate integrity. You can get pynacl from https://pypi.python.org/pypi/PyNaCl/"
 
 
@@ -113,8 +117,6 @@ def is_crypto_available():
     otherwise
   """
 
-  from stem.util import log
-
   try:
     from cryptography.utils import int_from_bytes, int_to_bytes
     from cryptography.hazmat.backends import default_backend
@@ -127,11 +129,56 @@ def is_crypto_available():
 
     return True
   except ImportError:
+    from stem.util import log
     log.log_once('stem.prereq.is_crypto_available', log.INFO, CRYPTO_UNAVAILABLE)
     return False
 
 
 @lru_cache()
+def is_zstd_available():
+  """
+  Checks if the `zstd module <https://pypi.python.org/pypi/zstandard>`_ is
+  available.
+
+  .. versionadded:: 1.7.0
+
+  :returns: **True** if we can use the zstd module and **False** otherwise
+  """
+
+  try:
+    # Unfortunately the zstandard module uses the same namespace as another
+    # zstd module (https://pypi.python.org/pypi/zstd), so we need to
+    # differentiate them.
+
+    import zstd
+    return hasattr(zstd, 'ZstdDecompressor')
+  except ImportError:
+    from stem.util import log
+    log.log_once('stem.prereq.is_zstd_available', log.INFO, ZSTD_UNAVAILABLE)
+    return False
+
+
+ at lru_cache()
+def is_lzma_available():
+  """
+  Checks if the `lzma module <https://docs.python.org/3/library/lzma.html>`_ is
+  available. This was added as a builtin in Python 3.3.
+
+  .. versionadded:: 1.7.0
+
+  :returns: **True** if we can use the lzma module and **False** otherwise
+  """
+
+  try:
+    import lzma
+    return True
+  except ImportError:
+    from stem.util import log
+    log.log_once('stem.prereq.is_lzma_available', log.INFO, LZMA_UNAVAILABLE)
+    return False
+
+
+ at lru_cache()
 def is_mock_available():
   """
   Checks if the mock module is available. In python 3.3 and up it is a builtin
diff --git a/test/settings.cfg b/test/settings.cfg
index f1cbd381..60224463 100644
--- a/test/settings.cfg
+++ b/test/settings.cfg
@@ -171,6 +171,7 @@ pyflakes.ignore stem/prereq.py => 'cryptography.hazmat.primitives.serialization.
 pyflakes.ignore stem/prereq.py => 'cryptography.hazmat.primitives.ciphers.modes' imported but unused
 pyflakes.ignore stem/prereq.py => 'cryptography.hazmat.primitives.ciphers.Cipher' imported but unused
 pyflakes.ignore stem/prereq.py => 'cryptography.hazmat.primitives.ciphers.algorithms' imported but unused
+pyflakes.ignore stem/prereq.py => 'lzma' imported but unused
 pyflakes.ignore stem/prereq.py => 'nacl.encoding' imported but unused
 pyflakes.ignore stem/prereq.py => 'nacl.signing' imported but unused
 pyflakes.ignore stem/interpreter/__init__.py => undefined name 'raw_input'
diff --git a/test/unit/descriptor/remote.py b/test/unit/descriptor/remote.py
index a02a1171..fd831e1e 100644
--- a/test/unit/descriptor/remote.py
+++ b/test/unit/descriptor/remote.py
@@ -207,20 +207,20 @@ class TestDescriptorDownloader(unittest.TestCase):
     self.assertEqual(TEST_RESOURCE, query.resource)
 
   def test_zstd_support_check(self):
-    with patch('stem.descriptor.remote.ZSTD_SUPPORTED', True):
+    with patch('stem.prereq.is_zstd_available', Mock(return_value = True)):
       query = stem.descriptor.remote.Query(TEST_RESOURCE, compression = Compression.ZSTD, start = False)
       self.assertEqual([Compression.ZSTD], query.compression)
 
-    with patch('stem.descriptor.remote.ZSTD_SUPPORTED', False):
+    with patch('stem.prereq.is_zstd_available', Mock(return_value = False)):
       query = stem.descriptor.remote.Query(TEST_RESOURCE, compression = Compression.ZSTD, start = False)
       self.assertEqual([Compression.PLAINTEXT], query.compression)
 
   def test_lzma_support_check(self):
-    with patch('stem.descriptor.remote.LZMA_SUPPORTED', True):
+    with patch('stem.prereq.is_lzma_available', Mock(return_value = True)):
       query = stem.descriptor.remote.Query(TEST_RESOURCE, compression = Compression.LZMA, start = False)
       self.assertEqual([Compression.LZMA], query.compression)
 
-    with patch('stem.descriptor.remote.LZMA_SUPPORTED', False):
+    with patch('stem.prereq.is_lzma_available', Mock(return_value = False)):
       query = stem.descriptor.remote.Query(TEST_RESOURCE, compression = Compression.LZMA, start = False)
       self.assertEqual([Compression.PLAINTEXT], query.compression)
 
@@ -260,7 +260,7 @@ class TestDescriptorDownloader(unittest.TestCase):
     Download a zstd compressed descriptor.
     """
 
-    if not stem.descriptor.remote.ZSTD_SUPPORTED:
+    if not stem.prereq.is_zstd_available():
       self.skipTest('(requires zstd module)')
       return
 
@@ -279,7 +279,7 @@ class TestDescriptorDownloader(unittest.TestCase):
     Download a lzma compressed descriptor.
     """
 
-    if not stem.descriptor.remote.LZMA_SUPPORTED:
+    if not stem.prereq.is_lzma_available():
       self.skipTest('(requires lzma module)')
       return
 





More information about the tor-commits mailing list