S7904: method for checking disk space is available

1 files changed, 21 insertions, 0 deletions

diff --git a/morphlib/util.py b/morphlib/util.py
index ffc84d58..1a347da0 100644
--- a/morphlib/util.py
+++ b/morphlib/util.py
@@ -236,3 +236,24 @@ def on_same_filesystem(path_a, path_b): # pragma: no cover
     """
     # TODO: return true if one path is a subvolume of the other on btrfs?
     return os.stat(path_a).st_dev == os.stat(path_b).st_dev
+
+def check_disk_available(tmp_path, tmp_min_size,
+                         cache_path, cache_min_size): # pragma: no cover
+    # if both are on the same filesystem, assume they share a storage pool,
+    # so the sum of the two sizes needs to be available
+    # TODO: if we need to do this on any more than 2 filesystems
+    #       extend it to take a [(path, min)] and do some arcane mathematics
+    #       to split it into groups that share a filesystem
+    #       hint: try list.sort and itertools.groupby
+    if on_same_filesystem(tmp_path, cache_path):
+        tmp_min_size = cache_min_size = tmp_min_size + cache_min_size
+    tmp_size, cache_size = map(get_bytes_free_in_path, (tmp_path, cache_path))
+    errors = []
+    for path, min in [(tmp_path, tmp_min_size), (cache_path, cache_min_size)]:
+        free = get_bytes_free_in_path(path)
+        if free < min:
+            errors.append('\t%(path)s requires %(min)d bytes free, '
+                          'has %(free)d' % locals())
+    if not errors:
+        return
+    raise morphlib.Error('Insufficient space on disk:\n' + '\n'.join(errors))