From 082958c20576deebc4c8a3cdcdfbcdb962e92bf9 Mon Sep 17 00:00:00 2001
From: Bryan Gardiner <bog@khumba.net>
Date: Sun, 2 Oct 2022 14:29:58 -0700
Subject: [PATCH] Add display of the change in closure disk size (issue #8).

---
 CHANGELOG.md |  2 ++
 src/nvd      | 63 +++++++++++++++++++++++++++++++++++++++++++++++++++-
 src/nvd.1    |  8 +++----
 3 files changed, 68 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index cc3b886..383efc0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## 0.1.3 (unreleased)
 
+- Add display of the change in closure disk size (issue #8).
+
 ## 0.1.2 (2021-11-05)
 
 - Added a flake.nix, thanks @dadada_.
diff --git a/src/nvd b/src/nvd
index 24eb64c..7dc1621 100755
--- a/src/nvd
+++ b/src/nvd
@@ -428,6 +428,55 @@ def print_package_list(
         ))
         count += 1
 
+def query_closure_disk_usage_bytes(target: Path) -> Optional[int]:
+    # If we don't add "./" to relative paths, then newer nix will interpret the
+    # argument as something related to flakes instead.
+    target_str = str(target)
+    if not target_str.startswith("/"):
+        target_str = "./" + target_str
+
+    try:
+        stdout = subprocess.run(
+                ["nix", "path-info", "--extra-experimental-features", "nix-command", "--closure-size", target_str],
+                stdout=PIPE,
+                text=True
+        ).stdout
+    except FileNotFoundError:
+        sys.stderr.write("nvd: Couldn't run 'nix path-info --closure-size'.\n")
+        return None
+
+    lines = stdout.rstrip("\n").split("\n")
+    if len(lines) != 1:
+        return None
+
+    words = lines[0].split()
+    return int(words[-1])
+
+BYTES_UNITS = ["KiB", "MiB", "GiB", "TiB", "PiB", "EiB"]
+
+def render_bytes(bytes_size: int) -> str:
+    scaled_size: Union[int, float] = bytes_size
+    scaled_suffix = "B"
+
+    neg = scaled_size < 0
+    if neg:
+        scaled_size = -scaled_size
+
+    for suffix in BYTES_UNITS:
+        if scaled_size >= 1024:
+            scaled_size /= 1024
+            scaled_suffix = suffix
+        else:
+            break
+
+    if neg:
+        scaled_size = -scaled_size
+
+    if isinstance(scaled_size, int):
+        return f"{scaled_size:+,d}{scaled_suffix}"  # B (bytes).
+    else:
+        return f"{scaled_size:+,.1f}{scaled_suffix}"  # KiB or higher.
+
 def run_list(*, root, only_selected, name_patterns):
     path = Path(root)
 
@@ -581,10 +630,22 @@ def run_diff(*, root1, root2):
     left_only_paths_count = len(left_closure_paths_set - right_closure_paths_set)
     right_only_paths_count = len(right_closure_paths_set - left_closure_paths_set)
 
+    left_closure_disk_usage_bytes: Optional[int] = \
+        query_closure_disk_usage_bytes(left_path)
+    right_closure_disk_usage_bytes: Optional[int] = \
+        query_closure_disk_usage_bytes(right_path)
+
+    if left_closure_disk_usage_bytes is None or right_closure_disk_usage_bytes is None:
+        diff_closure_disk_usage_str = ""
+    else:
+        diff_closure_disk_usage_str = ", disk usage " + \
+            render_bytes(right_closure_disk_usage_bytes - left_closure_disk_usage_bytes)
+
     print(
         f"Closure size: {len(left_closure_paths)} -> {len(right_closure_paths)} "
         f"({right_only_paths_count} paths added, {left_only_paths_count} paths removed, "
-        f"delta {right_only_paths_count - left_only_paths_count:+d})."
+        f"delta {right_only_paths_count - left_only_paths_count:+d}"
+        f"{diff_closure_disk_usage_str})."
     )
 
 def parse_args():
diff --git a/src/nvd.1 b/src/nvd.1
index fd0a83b..345b3bf 100644
--- a/src/nvd.1
+++ b/src/nvd.1
@@ -275,7 +275,7 @@ Added packages:
 [A+]  #5  neverball     1.6.0
 [A.]  #6  physfs        3.0.2
 [A.]  #7  smpeg-svn390  <none>
-Closure size: 2382 -> 2392 (21 paths added, 11 paths removed, delta +10).
+Closure size: 2382 -> 2392 (21 paths added, 11 paths removed, delta +10, disk usage +312.5MiB).
 .EE
 .RE
 .P
@@ -308,7 +308,7 @@ Added packages:
 [A+]  #1  gpsbabel        1.6.0
 [A+]  #2  gpxsee          7.31
 [A.]  #3  qttranslations  5.15.2
-Closure size: 2480 -> 2483 (44 paths added, 41 paths removed, delta +3).
+Closure size: 2480 -> 2483 (44 paths added, 41 paths removed, delta +3, disk usage +22.3MiB).
 .EE
 .RE
 .P
@@ -328,7 +328,7 @@ $ nvd diff /nix/var/nix/profiles/system-{43,44}-link
 >>> /nix/var/nix/profiles/system/system-44-link
 Version changes:
 [D.]  #1  nixos-system-unnamed  20.09.git.61092780ec5 -> 20.09.git.f79caa0b693
-Closure size: 2329 -> 2329 (16 paths added, 16 paths removed, delta +0).
+Closure size: 2329 -> 2329 (16 paths added, 16 paths removed, delta +0, disk usage +0B).
 .EE
 .RE
 .P
@@ -342,6 +342,6 @@ $ nvd diff /nix/var/nix/profiles/system-{23,24}-link
 <<< /nix/var/nix/profiles/system/system-23-link
 >>> /nix/var/nix/profiles/system/system-24-link
 No version or selection state changes.
-Closure size: 2191 -> 2191 (3 paths added, 3 paths removed, delta +0).
+Closure size: 2191 -> 2191 (3 paths added, 3 paths removed, delta +0, disk usage +0B).
 .EE
 .RE