Simplify setupext.download_or_cache.

anntzer · anntzer · commit a4ed88d1496b · 2019-10-16T23:08:25.000+02:00
Instead of repeatedly unwrapping and rewrapping the file contents in a
BytesIO that needs to be seek()ed, just pass around the bytes object
itself.
diff --git a/setupext.py b/setupext.py
@@ -36,19 +36,10 @@ def _get_xdg_cache_dir():
     return pathlib.Path(cache_dir, 'matplotlib')
 
 
-def get_fd_hash(fd):
-    """
-    Compute the sha256 hash of the bytes in a file-like
-    """
-    BLOCKSIZE = 1 << 16
+def _get_hash(data):
+    """Compute the sha256 hash of *data*."""
     hasher = hashlib.sha256()
-    old_pos = fd.tell()
-    fd.seek(0)
-    buf = fd.read(BLOCKSIZE)
-    while buf:
-        hasher.update(buf)
-        buf = fd.read(BLOCKSIZE)
-    fd.seek(old_pos)
+    hasher.update(data)
     return hasher.hexdigest()
 
 
@@ -59,64 +50,48 @@ def download_or_cache(url, sha):
     Parameters
     ----------
     url : str
-        The url to download
-
+        The url to download.
     sha : str
-        The sha256 of the file
+        The sha256 of the file.
 
     Returns
     -------
-    BytesIO
-        The file loaded into memory.
+    bytes
+        The file contents.
     """
     cache_dir = _get_xdg_cache_dir()
 
-    def get_from_cache(local_fn):
-        if cache_dir is None:
-            raise Exception("no cache dir")
-        buf = BytesIO((cache_dir / local_fn).read_bytes())
-        if get_fd_hash(buf) != sha:
-            return None
-        buf.seek(0)
-        return buf
-
-    def write_cache(local_fn, data):
-        if cache_dir is None:
-            raise Exception("no cache dir")
-        cache_dir.mkdir(parents=True, exist_ok=True)
-        old_pos = data.tell()
-        data.seek(0)
-        with open(cache_dir / local_fn, "xb") as fout:
-            fout.write(data.read())
-        data.seek(old_pos)
-
-    try:
-        return get_from_cache(sha)
-    except Exception:
-        pass
+    if cache_dir is not None:  # Try to read from cache.
+        try:
+            data = (cache_dir / sha).read_bytes()
+        except IOError:
+            pass
+        else:
+            if _get_hash(data) == sha:
+                return data
 
     # jQueryUI's website blocks direct downloads from urllib.request's
     # default User-Agent, but not (for example) wget; so I don't feel too
     # bad passing in an empty User-Agent.
     with urllib.request.urlopen(
             Request(url, headers={"User-Agent": ""})) as req:
-        file_contents = BytesIO(req.read())
-        file_contents.seek(0)
-
-    file_sha = get_fd_hash(file_contents)
+        data = req.read()
 
+    file_sha = _get_hash(data)
     if file_sha != sha:
         raise Exception(
             f"The download file does not match the expected sha.  {url} was "
             f"expected to have {sha} but it had {file_sha}")
 
-    try:
-        write_cache(sha, file_contents)
-    except Exception:
-        pass
+    if cache_dir is not None:  # Try to cache the downloaded file.
+        try:
+            cache_dir.mkdir(parents=True, exist_ok=True)
+            with open(cache_dir / sha, "xb") as fout:
+                fout.write(data)
+        except IOError:
+            pass
 
-    file_contents.seek(0)
-    return file_contents
+    return data
 
 
 # SHA256 hashes of the FreeType tarballs
@@ -192,16 +167,6 @@ def print_status(package, status):
                             subsequent_indent=indent))
 
 
-def get_buffer_hash(fd):
-    BLOCKSIZE = 1 << 16
-    hasher = hashlib.sha256()
-    buf = fd.read(BLOCKSIZE)
-    while buf:
-        hasher.update(buf)
-        buf = fd.read(BLOCKSIZE)
-    return hasher.hexdigest()
-
-
 def deplib(libname):
     if sys.platform != 'win32':
         return libname
@@ -545,18 +510,13 @@ def do_custom_build(self):
         if not src_path.exists():
             os.makedirs('build', exist_ok=True)
 
-            url_fmts = [
-                ('https://downloads.sourceforge.net/project/freetype'
-                 '/freetype2/{version}/{tarball}'),
-                ('https://download.savannah.gnu.org/releases/freetype'
-                 '/{tarball}')
-            ]
             tarball = f'freetype-{LOCAL_FREETYPE_VERSION}.tar.gz'
-
             target_urls = [
-                url_fmt.format(version=LOCAL_FREETYPE_VERSION,
-                               tarball=tarball)
-                for url_fmt in url_fmts]
+                (f'https://downloads.sourceforge.net/project/freetype'
+                 f'/freetype2/{LOCAL_FREETYPE_VERSION}/{tarball}'),
+                (f'https://download.savannah.gnu.org/releases/freetype'
+                 f'/{tarball}')
+            ]
 
             for tarball_url in target_urls:
                 try:
@@ -572,10 +532,8 @@ def do_custom_build(self):
                     f"top-level of the source repository.")
 
             print(f"Extracting {tarball}")
-            # just to be sure
-            tar_contents.seek(0)
-            with tarfile.open(tarball, mode="r:gz",
-                              fileobj=tar_contents) as tgz:
+            with tarfile.open(fileobj=BytesIO(tar_contents),
+                              mode="r:gz") as tgz:
                 tgz.extractall("build")
 
         print(f"Building freetype in {src_path}")
