From 69ab9a34468e0d96d8d429a9aad35a14c1170b5b Mon Sep 17 00:00:00 2001
From: yut23 <yut23@gvljohnsons.com>
Date: Wed, 24 Apr 2024 14:56:03 -0400
Subject: [PATCH] Add formatting based on the timeit CLI interface

---
 src/yut23_utils/timing.py |  6 ++++++
 tests/test_timing.py      | 22 ++++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/src/yut23_utils/timing.py b/src/yut23_utils/timing.py
index dc030a4..18abf49 100644
--- a/src/yut23_utils/timing.py
+++ b/src/yut23_utils/timing.py
@@ -42,6 +42,7 @@ def _format_time(
 class TimingFormat(enum.Enum):
     IPYTHON = enum.auto()
     HYPERFINE = enum.auto()
+    TIMEIT = enum.auto()
 
 
 @dataclass(frozen=True)
@@ -113,6 +114,11 @@ def time_str(timespan: float) -> str:
                 f" (mean ± std. dev. of {label_count(self.repeat, 'loop')},"
                 f" {loop_str} each)"
             )
+        if fmt is TimingFormat.TIMEIT:
+            return (
+                f"{loop_str}, best of {self.repeat}:"
+                f" {_format_time(self.min)} per loop"
+            )
         raise AssertionError
 
 
diff --git a/tests/test_timing.py b/tests/test_timing.py
index c204ac4..f885759 100644
--- a/tests/test_timing.py
+++ b/tests/test_timing.py
@@ -86,3 +86,25 @@ def test_pretty_ipython_single(self):
         assert info.pretty(TimingFormat.IPYTHON) == (
             "3.14 s ± nan s per loop (mean ± std. dev. of 1 loop, 2 loops each)"
         )
+
+    def test_pretty_timeit(self):
+        times = (1.23, 3.21, 2.75, 2.53)
+
+        info = TimingInfo(times, 1)
+        assert info.pretty(TimingFormat.TIMEIT) == "1 loop, best of 4: 1.23 s per loop"
+
+        info = TimingInfo(times, 3)
+        assert info.pretty(TimingFormat.TIMEIT) == "3 loops, best of 4: 1.23 s per loop"
+
+    def test_pretty_timeit_single(self):
+        times = (3.14159,)
+
+        info = TimingInfo(times, 1)
+        assert info.pretty(TimingFormat.TIMEIT) == (
+            "1 loop, best of 1: 3.14 s per loop"
+        )
+
+        info = TimingInfo(times, 2)
+        assert info.pretty(TimingFormat.TIMEIT) == (
+            "2 loops, best of 1: 3.14 s per loop"
+        )