Progress bar for webviz store/assets (#251)

equinor · Jul 29, 2020 · d1410b4 · d1410b4
1 parent 5f06334
commit d1410b4
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 48 deletions.
diff --git a/setup.py b/setup.py
@@ -44,6 +44,7 @@
         "pandas>=0.24",
         "pyarrow>=0.16",
         "pyyaml>=5.1",
+        "tqdm>=4.8",
         "typing-extensions>=3.7",  # Needed on Python < 3.8
         "webviz-core-components>=0.0.19",
     ],

diff --git a/webviz_config/webviz_assets.py b/webviz_config/webviz_assets.py
@@ -4,11 +4,10 @@
 import pathlib
 from typing import Optional
 
+from tqdm import tqdm
 from dash import Dash
 import flask
 
-from .utils import terminal_colors
-
 
 class WebvizAssets:
     """Dash applications by default host static resources from a folder called
@@ -108,21 +107,13 @@ def make_portable(self, asset_folder: str) -> None:
         """Copy over all added assets to the given folder (asset_folder).
         """
 
-        for counter, (assigned_id, filename) in enumerate(self._assets.items()):
-            print(
-                f"{terminal_colors.PURPLE} Copying over {filename} {terminal_colors.END}",
-                end="",
-                flush=True,
-            )
-
+        for assigned_id, filename in tqdm(
+            self._assets.items(),
+            bar_format="{l_bar} {bar} | Copied {n_fmt}/{total_fmt}",
+        ):
+            tqdm.write(f"Copying over {filename}")
             shutil.copyfile(filename, os.path.join(asset_folder, assigned_id))
 
-            print(
-                f"{terminal_colors.PURPLE}{terminal_colors.BOLD} "
-                f"[\u2713] Copied ({counter + 1}/{len(self._assets)})"
-                f"{terminal_colors.END}"
-            )
-
     def _generate_id(self, filename: str) -> str:
         """From the filename, create a safe resource id not already present
         """

diff --git a/webviz_config/webviz_store.py b/webviz_config/webviz_store.py
@@ -10,8 +10,7 @@
 from typing import Callable, List, Union, Any
 
 import pandas as pd
-
-from .utils import terminal_colors
+from tqdm import tqdm
 
 
 class WebvizStorage:
@@ -162,37 +161,28 @@ def build_store(self) -> None:
             len(calls) for calls in self.storage_function_argvalues.values()
         )
 
-        counter = 0
-        for func in self.storage_functions:
-            for argtuples in self.storage_function_argvalues[func].values():
-                kwargs = dict(argtuples)
-
-                print(
-                    f"{terminal_colors.PURPLE}"
-                    f" Running {WebvizStorage.string(func, kwargs)}"
-                    f"{terminal_colors.END}",
-                    end="",
-                    flush=True,
-                )
-
-                output = func(**kwargs)
-                path = self._unique_path(func, argtuples)
-
-                if isinstance(output, pd.DataFrame):
-                    output.to_parquet(f"{path}.parquet")
-                elif isinstance(output, pathlib.Path):
-                    shutil.copy(output, f"{path}{output.suffix}")
-                elif isinstance(output, io.BytesIO):
-                    pathlib.Path(path).write_bytes(output.getvalue())
-                else:
-                    raise ValueError(f"Unknown return type {type(output)}")
-
-                counter += 1
-                print(
-                    f"{terminal_colors.PURPLE}{terminal_colors.BOLD}"
-                    f"[\u2713] Saved ({counter}/{total_calls})"
-                    f"{terminal_colors.END}"
-                )
+        with tqdm(
+            total=total_calls, bar_format="{l_bar} {bar} | Saved {n_fmt}/{total_fmt}"
+        ) as progress_bar:
+            for func in self.storage_functions:
+                if self.storage_function_argvalues[func]:
+                    progress_bar.write(
+                        f"Storing output of {func.__module__}.{func.__name__}"
+                    )
+                for argtuples in self.storage_function_argvalues[func].values():
+                    output = func(**dict(argtuples))
+                    path = self._unique_path(func, argtuples)
+
+                    if isinstance(output, pd.DataFrame):
+                        output.to_parquet(f"{path}.parquet")
+                    elif isinstance(output, pathlib.Path):
+                        shutil.copy(output, f"{path}{output.suffix}")
+                    elif isinstance(output, io.BytesIO):
+                        pathlib.Path(path).write_bytes(output.getvalue())
+                    else:
+                        raise ValueError(f"Unknown return type {type(output)}")
+
+                    progress_bar.update()
 
 
 def webvizstore(func: Callable) -> Callable: