Cleanup new profiler.

J08nY · Jul 15, 2024 · 73e3bf1 · 73e3bf1
1 parent d71a021
commit 73e3bf1
Showing 1 changed file with 15 additions and 2 deletions.
diff --git a/test/utils.py b/test/utils.py
@@ -11,6 +11,10 @@
 
 
 class RawTimer:
+    start: int
+    end: int
+    duration: float
+
     def __enter__(self):
         self.start = time.perf_counter_ns()
 
@@ -20,8 +24,17 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
 
 class Profiler:
-    def __init__(self, prof_type: Union[Literal["py"], Literal["c"], Literal["raw"]], output_directory: str, benchmark_name: str):
-        self._prof: Union[PyProfiler, cProfiler, RawTimer] = {"py": PyProfiler, "c": cProfiler, "raw": RawTimer}[prof_type]()
+    def __init__(
+        self,
+        prof_type: Union[Literal["py"], Literal["c"], Literal["raw"]],
+        output_directory: str,
+        benchmark_name: str,
+    ):
+        self._prof: Union[PyProfiler, cProfiler, RawTimer] = {
+            "py": PyProfiler,
+            "c": cProfiler,
+            "raw": RawTimer,
+        }[prof_type]()
         self._prof_type: Union[Literal["py"], Literal["c"], Literal["raw"]] = prof_type
         self._root_frame = None
         self._state = "out"