Update on "Simple benchmark of ufmt_file and ufmt_paths"

[ghstack-poisoned]

ufmt/tests/benchmark.py

@@ -12,78 +12,69 @@
 ROOT = Path(__file__).parent.parent.parent
 
 
-class Timer:
-    def __init__(self, name: str) -> None:
-        self.name = name
+class Benchmark:
+    def __init__(self, target: int = 5) -> None:
+        self.target = target
+        self.last_time: int = 0
+        self.title: str = ""
         self.totals: List[int] = []
 
     @classmethod
     def fields(self) -> str:
-        headline = f"{'name':^40}  {'min':^10}  {'mean':^10}  {'max':^10}"
+        headline = f"{'benchmark':^40}  {'min':^10}  {'mean':^10}  {'max':^10}"
         underline = "-" * len(headline)
         return f"{headline}\n{underline}"
 
-    def __str__(self) -> str:
-        short = min(self.totals)
-        long = max(self.totals)
-        avg = sum(self.totals) // len(self.totals)
+    def results(self) -> str:
+        totals = self.totals
+        short = min(totals)
+        long = max(totals)
+        avg = sum(totals) // len(totals)
         fields = "  ".join(f"{value // 1000:>7} µs" for value in (short, avg, long))
-        return f"{self.name + ':':<40}  {fields}"
+        return f"{self.title + ':':<40}  {fields}"
 
     def __enter__(self) -> Self:
-        self.before = time.monotonic_ns()
+        print()
+        print(self.fields())
         return self
 
     def __exit__(self, *args: Any) -> None:
-        after = time.monotonic_ns()
-        self.totals.append(after - self.before)
+        print()
 
+    def __call__(self, name: str) -> Self:
+        self.title = name
+        return self
 
-def benchmark() -> None:
-    print("starting benchmark...")
+    def __iter__(self) -> Self:
+        self.last_time = 0
+        self.totals = []
+        return self
+
+    def __next__(self) -> None:
+        now = time.monotonic_ns()
+        if self.last_time > 0:
+            self.totals += [now - self.last_time]
+        if len(self.totals) >= self.target:
+            print(self.results())
+            self.totals = []
+            self.title = ""
 
+            raise StopIteration
+        self.last_time = time.monotonic_ns()
+
+
+def benchmark() -> None:
     ufmt_dir = ROOT / "ufmt"
     ufmt_core = ufmt_dir / "core.py"
     assert ufmt_dir.is_dir(), f"{ufmt_dir} not found, must run benchmark from repo"
 
-    print()
-    print(Timer.fields())
+    with Benchmark() as benchmark:
 
-    timer = Timer("ufmt_file")
-    for _ in range(5):
-        with timer:
+        for _ in benchmark("ufmt_file"):
             ufmt_file(ufmt_core, dry_run=True)
-    print(timer)
-
-    timer = Timer("ufmt_file, diff=True")
-    for _ in range(5):
-        with timer:
-            ufmt_file(ufmt_core, dry_run=True, diff=True)
-    print(timer)
-
-    timer = Timer("ufmt_file, return_content=True")
-    for _ in range(5):
-        with timer:
-            ufmt_file(ufmt_core, dry_run=True, return_content=True)
-    print(timer)
-
-    timer = Timer("ufmt_paths")
-    for _ in range(5):
-        with timer:
+
+        for _ in benchmark("ufmt_paths"):
             list(ufmt_paths([ufmt_dir], dry_run=True))
-    print(timer)
-
-    timer = Timer("ufmt_paths, diff=True")
-    for _ in range(5):
-        with timer:
-            list(ufmt_paths([ufmt_dir], dry_run=True, diff=True))
-    print(timer)
-
-    timer = Timer("ufmt_paths, return_content=True")
-    for _ in range(5):
-        with timer:
-            list(ufmt_paths([ufmt_dir], dry_run=True, diff=True))
-    print(timer)
 
 
 if __name__ == "__main__":