stdlib: Add progress bars for long functions

This adds a progress bar for downloading large files and computing
md5sums on large files.

Change-Id: Iddc9faf61e861837cc1e2e3b3dbdbeebd6ccf529
Reviewed-on: https://gem5-review.googlesource.com/c/public/gem5/+/67472
Tested-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Melissa Jost <melissakjost@gmail.com>
Maintainer: Jason Lowe-Power <power.jg@gmail.com>

optional-requirements.txt

@@ -0,0 +1 @@
+tqdm==4.64.1

src/python/SConscript

@@ -267,6 +267,7 @@ PySource('gem5.resources', 'gem5/resources/elfie.py')
 PySource('gem5.utils', 'gem5/utils/__init__.py')
 PySource('gem5.utils', 'gem5/utils/filelock.py')
 PySource('gem5.utils', 'gem5/utils/override.py')
+PySource('gem5.utils', 'gem5/utils/progress_bar.py')
 PySource('gem5.utils', 'gem5/utils/requires.py')
 PySource('gem5.utils.multiprocessing',
     'gem5/utils/multiprocessing/__init__.py')

src/python/gem5/resources/downloader.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2021 The Regents of the University of California
+# Copyright (c) 2021-2023 The Regents of the University of California
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -42,6 +42,7 @@ from urllib.error import HTTPError
 from typing import List, Dict, Set, Optional
 
 from .md5_utils import md5_file, md5_dir
+from ..utils.progress_bar import tqdm, progress_hook
 
 from ..utils.filelock import FileLock
 
@@ -286,10 +287,26 @@ def _download(url: str, download_to: str, max_attempts: int = 6) -> None:
                 # get the file as a bytes blob
                 request = urllib.request.Request(url)
                 with urllib.request.urlopen(request, context=ctx) as fr:
-                    with open(download_to, "wb") as fw:
-                        fw.write(fr.read())
+                    with tqdm.wrapattr(
+                        open(download_to, "wb"),
+                        "write",
+                        miniters=1,
+                        desc="Downloading {download_to}",
+                        total=getattr(fr, "length", None),
+                    ) as fw:
+                        for chunk in fr:
+                            fw.write(chunk)
             else:
-                urllib.request.urlretrieve(url, download_to)
+                with tqdm(
+                    unit="B",
+                    unit_scale=True,
+                    unit_divisor=1024,
+                    miniters=1,
+                    desc=f"Downloading {download_to}",
+                ) as t:
+                    urllib.request.urlretrieve(
+                        url, download_to, reporthook=progress_hook(t)
+                    )
             return
         except HTTPError as e:
             # If the error code retrieved is retryable, we retry using a

src/python/gem5/resources/md5_utils.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2022 The Regents of the University of California
+# Copyright (c) 2022-2023 The Regents of the University of California
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -31,7 +31,22 @@ from _hashlib import HASH as Hash
 
 def _md5_update_from_file(filename: Path, hash: Hash) -> Hash:
     assert filename.is_file()
-    with open(str(filename), "rb") as f:
+
+    if filename.stat().st_size < 1024 * 1024 * 100:
+        from ..utils.progress_bar import FakeTQDM
+
+        # if the file is less than 100MB, no need to show a progress bar.
+        tqdm = FakeTQDM()
+    else:
+        from ..utils.progress_bar import tqdm
+
+    with tqdm.wrapattr(
+        open(str(filename), "rb"),
+        "read",
+        miniters=1,
+        desc=f"Computing md5sum on {filename}",
+        total=filename.stat().st_size,
+    ) as f:
         for chunk in iter(lambda: f.read(4096), b""):
             hash.update(chunk)
     return hash

src/python/gem5/utils/progress_bar.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2023 The Regents of the University of California
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are
+# met: redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer;
+# redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution;
+# neither the name of the copyright holders nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+class FakeTQDM:
+    """This is a fake wrapper so that the tqdm calls work whether or not it
+    has been installed.
+    """
+
+    def __call__(*args, **kwargs):
+        if args:
+            return args[0]
+        return kwargs.get("iterable", None)
+
+    def wrapattr(self, *args, **kwargs):
+        if args:
+            return args[0]
+        return kwargs.get("iterable", None)
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+
+try:
+    from tqdm.auto import tqdm
+
+    _have_tqdm = True
+except ImportError:
+    tqdm = FakeTQDM()
+    _have_tqdm = False
+
+# Hook for the progress bar
+def progress_hook(t):
+    if not _have_tqdm:
+        # Takes 3 arguments
+        return lambda a, b, c: None
+
+    last_b = [0]
+
+    def update_to(b=1, bsize=1, tsize=None):
+        if tsize not in (None, -1):
+            t.total = tsize
+        displayed = t.update((b - last_b[0]) * bsize)
+        last_b[0] = b
+        return displayed
+
+    return update_to
+
+
+__all__ = [tqdm, progress_hook, FakeTQDM]