Adds report output to repartition script to help visualise its output

Author: zooba

scripts/repartition-index.py

@@ -1,10 +1,9 @@
+import contextlib
 import json
 import re
 import sys
 
-from collections import OrderedDict
 from pathlib import Path
-from urllib.request import Request, urlopen
 
 REPO = Path(__file__).absolute().parent.parent
 sys.path.append(str(REPO / "src"))
@@ -30,6 +29,7 @@ def usage():
     print("  -t/--tag TAG       Include only the specified tags (comma-separated)")
     print("  -r/--range RANGE   Include only the specified range (comma-separated)")
     print("  --latest-micro     Include only the latest x.y.z version")
+    print("  --report           Write plain-text summary report")
     print()
     print("An output of 'nul' is permitted to drop entries.")
     print("Providing the same inputs and outputs is permitted, as all inputs are read")
@@ -96,6 +96,7 @@ def __init__(self):
         self.tag_or_range = None
         self._expect_tag_or_range = False
         self.latest_micro = False
+        self.report = False
 
     def add_arg(self, arg):
         if arg[:1] != "-":
@@ -121,9 +122,16 @@ def add_arg(self, arg):
         if arg == "--latest-micro":
             self.latest_micro = True
             return False
+        if arg == "--report":
+            self.report = True
+            return False
         raise ValueError("Unknown argument: " + arg)
 
     def execute(self, versions, context):
+        if self.report:
+            if self.target != "nul":
+                context.setdefault("reports", []).append(self.target)
+            return
         written = context.setdefault("written", set())
         written_now = set()
         outputs = context.setdefault("outputs", {})
@@ -176,19 +184,56 @@ def add_arg(self, arg):
             return False
         raise ValueError("Unknown argument: " + arg)
 
+    @contextlib.contextmanager
+    def open(self, file):
+        file = Path(file)
+        if file.match("nul"):
+            import io
+            yield io.StringIO()
+        elif file.match("stdout"):
+            yield sys.stdout
+        else:
+            with open(file, "w", encoding="utf-8") as f:
+                yield f
+
+    def st_size(self, file):
+        file = Path(file)
+        if file.match("nul"):
+            return "no data written"
+        if file.match("stdout"):
+            return "n/a"
+        return f"{Path(file).stat().st_size} bytes"
+
     def execute(self, versions, context):
         outputs = context.get("outputs") or {}
         output_order = context.get("output_order", [])
+        report_data = {}
         for target, next_target in zip(output_order, [*output_order[1:], None]):
             data = {
                 "versions": outputs[target]
             }
             if next_target:
                 data["next"] = next_target
-            with open(target, "w", encoding="utf-8") as f:
+            for i in outputs[target]:
+                report_data.setdefault(target, {}).setdefault(i["sort-version"].casefold(), []).append(i)
+            with self.open(target) as f:
                 json.dump(data, f, indent=self.indent)
             print("Wrote {} ({} entries, {} bytes)".format(
-                target, len(data["versions"]), Path(target).stat().st_size
+                target, len(data["versions"]), self.st_size(target)
+            ))
+
+        reports = context.get("reports", [])
+        for target in reports:
+            with self.open(target) as f:
+                for output_target in output_order:
+                    print("Written to", output_target, file=f)
+                    data = report_data[output_target]
+                    for key in data:
+                        ids = ", ".join(i["id"] for i in data[key])
+                        print("{}: {}".format(key, ids), file=f)
+                    print(file=f)
+            print("Wrote {} ({} bytes)".format(
+                target, self.st_size(target)
             ))
 
 
@@ -203,10 +248,13 @@ def parse_cli(args):
             print("Using equivalent of: --pre --latest-micro -r >=3.11.0 index-windows.json")
             print("                     --pre -r >=3.11.0 index-windows-recent.json")
             print("                     index-windows-legacy.json")
-            plan_split = [SplitToFile(), SplitToFile(), SplitToFile()]
+            print("                     --report index-windows.txt")
+            plan_split = [SplitToFile(), SplitToFile(), SplitToFile(), SplitToFile()]
             plan_split[0].target = "index-windows.json"
             plan_split[1].target = "index-windows-recent.json"
             plan_split[2].target = "index-windows-legacy.json"
+            plan_split[3].target = "stdout"
+            plan_split[3].report = True
             plan_split[0].pre = plan_split[1].pre = plan_split[2].pre = True
             plan_split[0].latest_micro = True
             plan_split[0].tag_or_range = tag_or_range(">=3.11.0")
@@ -246,4 +294,3 @@ def parse_cli(args):
     CONTEXT = {}
     for p in plan:
         p.execute(VERSIONS, CONTEXT)
-