Add argument to 'clean' command to remove reports files

lib/ramble/ramble/cmd/clean.py

@@ -11,9 +11,13 @@
 import os
 import shutil
 
+import llnl.util.tty as tty
+import spack.util.spack_yaml as syaml
+
 import ramble.caches
 import ramble.config
 import ramble.repository
+import ramble.reports
 import ramble.stage
 from ramble.util.logger import logger
 from ramble.paths import lib_path, var_path
@@ -43,12 +47,18 @@ def setup_parser(subparser):
         action="store_true",
         help="remove .pyc, .pyo files and __pycache__ folders",
     )
+    subparser.add_argument(
+        "-r",
+        "--reports",
+        action="store_true",
+        help="remove pdf and image files generated by ramble reports",
+    )
     subparser.add_argument("-a", "--all", action=AllClean, help="equivalent to -dmp", nargs=0)
 
 
 def clean(parser, args):
     # If nothing was set, activate the default
-    if not any([args.downloads, args.misc_cache, args.python_cache]):
+    if not any([args.downloads, args.misc_cache, args.python_cache, args.reports]):
         args.downloads = True
 
     if args.downloads:
@@ -63,6 +73,10 @@ def clean(parser, args):
         logger.msg("Removing python cache files")
         remove_python_caches()
 
+    if args.reports:
+        logger.msg("Removing pdf and image files generated by ramble reports")
+        remove_reports_files()
+
 
 def remove_python_caches():
     logger.msg("Removing python cache files")
@@ -78,3 +92,44 @@ def remove_python_caches():
                     dname = os.path.join(root, d)
                     logger.debug(f"Removing {dname}")
                     shutil.rmtree(dname)
+
+
+def remove_reports_files():
+    reports_path = ramble.reports.get_reports_path()
+    if reports_path:
+
+        answer = tty.get_yes_or_no(
+            f"Really remove all reports and images from {reports_path}?",
+            default=False,
+        )
+        if not answer:
+            logger.die("Will not remove any files")
+
+        for root, _, files in os.walk(reports_path, topdown=False):
+            inventory_file = os.path.join(root, ramble.reports.INVENTORY_FILENAME)
+
+            try:
+                with open(inventory_file) as f:
+                    inventory = syaml.load(f)
+            except FileNotFoundError:
+                continue
+
+            if inventory:
+                for inv_file in inventory["files"]:
+                    if inv_file in files:
+                        fname = os.path.join(root, inv_file)
+                        logger.debug(f"Removing {fname}")
+                        os.remove(fname)
+                logger.debug(f"Removing {inventory_file}")
+                os.remove(inventory_file)
+
+            if not os.listdir(root):
+                logger.debug(f"Removing empty directory {root}")
+                os.rmdir(root)
+
+        # Clean up symlinks in root dir
+        for item in os.listdir(reports_path):
+            item_path = os.path.join(reports_path, item)
+            if os.path.islink(item_path):
+                logger.debug(f"Removing {item_path}")
+                os.remove(item_path)

lib/ramble/ramble/reports.py

@@ -13,6 +13,7 @@
 import re
 
 import llnl.util.filesystem as fs
+import spack.util.spack_yaml as syaml
 
 import ramble.cmd.results
 import ramble.cmd.workspace
@@ -56,6 +57,8 @@ class ReportVars(Enum):
     "origin_type": ReportVars.FOM_ORIGIN_TYPE.value,
 }
 
+INVENTORY_FILENAME = "inventory.yaml"
+
 
 def to_numeric_if_possible(series):
     """Try to convert a Pandas series to numeric, or return the series unchanged."""
@@ -262,6 +265,7 @@ def __init__(self, spec, normalize, report_dir_path, results_df, logx, logy, spl
         self.normalize = normalize
         self.spec = spec
         self.report_dir_path = report_dir_path
+        self.inventory_path = os.path.join(report_dir_path, INVENTORY_FILENAME)
         self.figsize = [12, 8]
 
         self.results_df = results_df
@@ -326,6 +330,23 @@ def add_minmax_data(self, selected_data, min_data, max_data, scale_var):
                 from_col=ReportVars.FOM_VALUE_MAX.value,
             )
 
+    def add_to_inventory(self, filename):
+        """Adds a filename to the inventory.
+
+        Args:
+            filename: filename to add to inventory.
+        """
+        try:
+            with open(self.inventory_path) as f:
+                inventory = syaml.load(f)
+        except FileNotFoundError:
+            inventory = {"files": []}
+
+        inventory["files"].append(filename)
+
+        with open(self.inventory_path, "w+") as f:
+            syaml.dump(inventory, stream=f)
+
     def draw(self, perf_measure, scale_var, series, pdf_report, y_label=None):
         series_data = self.output_df.query(f'series == "{series}"').copy()
 
@@ -431,6 +452,7 @@ def validate_spec(self, chart_spec):
     def write(self, fig, filename, pdf_report):
         filename = filename.replace(" ", "-")
         plt.savefig(os.path.join(self.report_dir_path, filename))
+        self.add_to_inventory(filename)
         pdf_report.savefig(fig)
         plt.close(fig)
 
@@ -889,9 +911,11 @@ def make_report(results_df, ws_name, args):
     plot = plot_factory.create_plot_generator(args, report_dir_path, results_df)
     plot_type = plot.plot_type
 
-    pdf_path = os.path.join(report_dir_path, f"{report_name}.{plot_type}.pdf")
+    pdf_filename = f"{report_name}.{plot_type}.pdf"
+    pdf_path = os.path.join(report_dir_path, pdf_filename)
     with PdfPages(pdf_path) as pdf_report:
         plot.generate_plot_data(pdf_report)
+        plot.add_to_inventory(pdf_filename)
 
     if os.path.isfile(pdf_path):
 

share/ramble/ramble-completion.bash

@@ -281,7 +281,7 @@ _ramble_attributes() {
 }
 
 _ramble_clean() {
-    RAMBLE_COMPREPLY="-h --help -d --downloads -m --misc-cache -p --python-cache -a --all"
+    RAMBLE_COMPREPLY="-h --help -d --downloads -m --misc-cache -p --python-cache -r --reports -a --all"
 }
 
 _ramble_commands() {