tools to check table disc usage

Author: samuelbray32

src/spyglass/utils/dj_helper_fn.py

@@ -559,6 +559,38 @@ def make_file_obj_id_unique(nwb_path: str):
     return new_id
 
 
+def _quick_get_analysis_path(file):
+    """Get the absolute path to an analysis file on disc without integrity checks.
+    For use when scanning large number of files.
+
+    Parameters
+    ----------
+    file : str
+        Name of the analysis file to get the path for.
+    Returns
+    -------
+    str
+        Absolute path to the analysis file. Returns None if the file is not found locally.
+    """
+    from spyglass.common import AnalysisNwbfile
+    from spyglass.settings import sg_config
+
+    analysis_dir = sg_config.analysis_dir
+    if os.path.exists(path := f"{analysis_dir}/{file}"):
+        return path
+
+    folder = "_".join(file.split("_")[:-1])
+    path = f"{analysis_dir}/{folder}/{file}"
+    if os.path.exists(path):
+        return path
+    path = AnalysisNwbfile().get_abs_path(file)
+    if os.path.exists(path):
+        return path
+
+    print(f"File {file} not found in {analysis_dir}")
+    return None
+
+
 def populate_pass_function(value):
     """Pass function for parallel populate.
 
@@ -612,3 +644,4 @@ def str_to_bool(value) -> bool:
     if not value:
         return False
     return str(value).lower() in ("y", "yes", "t", "true", "on", "1")
+    return str(value).lower() in ("y", "yes", "t", "true", "on", "1")

src/spyglass/utils/dj_mixin.py

@@ -1,9 +1,11 @@
+import os
 from contextlib import nullcontext
 from functools import cached_property
 from time import time
 from typing import List
 
 import datajoint as dj
+import humanize
 from datajoint.condition import make_condition
 from datajoint.errors import DataJointError
 from datajoint.expression import QueryExpression
@@ -16,6 +18,7 @@
 from spyglass.utils.database_settings import SHARED_MODULES
 from spyglass.utils.dj_helper_fn import (
     NonDaemonPool,
+    _quick_get_analysis_path,
     ensure_names,
     fetch_nwb,
     get_nwb_table,
@@ -863,6 +866,31 @@ def check_threads(self, detailed=False, all_threads=False) -> DataFrame:
 
         return df
 
+    # --------------------------- Check disc usage ------------------------------
+    def get_table_storage_usage(self):
+        """Total size of all analysis files in the table.
+        Uses the analysis_file_name field to find the file paths and sum their
+        sizes.
+
+        Returns
+        -------
+        tuple
+            (human-readable string, total size in bytes)
+        """
+        if "analysis_file_name" not in self.heading.names:
+            logger.warning(
+                f"{self.full_table_name} does not have an analysis_file_name field."
+            )
+            return "0 Mib", 0
+        file_names = self.fetch("analysis_file_name")
+        file_paths = [
+            _quick_get_analysis_path(file_name) for file_name in file_names
+        ]
+        file_paths = [path for path in file_paths if path is not None]
+        file_sizes = [os.stat(path).st_size for path in file_paths]
+        total_size = sum(file_sizes)
+        return humanize.naturalsize(total_size, binary=True), total_size
+
 
 class SpyglassMixinPart(SpyglassMixin, dj.Part):
     """