1 files changed, 24 insertions, 9 deletions

diff --git a/test/utils.py b/test/utils.py
index 0939528..1b00ef7 100644
--- a/test/utils.py
+++ b/test/utils.py
@@ -1,19 +1,30 @@
 import pstats
 import sys
+import time
 
 from pathlib import Path
 from subprocess import run, PIPE, DEVNULL
+from typing import Union, Literal
 
 from pyinstrument import Profiler as PyProfiler
 from cProfile import Profile as cProfiler
 
 
+class RawTimer:
+    def __enter__(self):
+        self.start = time.perf_counter_ns()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.end = time.perf_counter_ns()
+        self.duration = (self.end - self.start) / 1e9
+
+
 class Profiler:
-    def __init__(self, prof_type, output_directory, benchmark_name):
-        self._prof = PyProfiler() if prof_type == "py" else cProfiler()
-        self._prof_type = 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 = None
+        self._state = "out"
         self._output_directory = output_directory
         self._benchmark_name = benchmark_name
 
@@ -25,7 +36,7 @@ class Profiler:
     def __exit__(self, exc_type, exc_val, exc_tb):
         self._prof.__exit__(exc_type, exc_val, exc_tb)
         if self._prof_type == "py":
-            self._root_frame = self._prof.last_session.root_frame()
+            self._root_frame = self._prof.last_session.root_frame()  # type: ignore
         self._state = "out"
         self.output()
         self.save()
@@ -62,14 +73,18 @@ class Profiler:
         if self._state != "out":
             raise ValueError
         if self._prof_type == "py":
-            print(self._prof.output_text(unicode=True, color=True))
-        else:
-            self._prof.print_stats("cumtime")
+            print(self._prof.output_text(unicode=True, color=True))  # type: ignore
+        elif self._prof_type == "c":
+            self._prof.print_stats("cumtime")  # type: ignore
+        elif self._prof_type == "raw":
+            print(f"{self._prof.duration:.4} s")  # type: ignore
 
     def get_time(self) -> float:
         if self._state != "out":
             raise ValueError
         if self._prof_type == "py":
             return self._root_frame.time  # type: ignore
-        else:
+        elif self._prof_type == "c":
             return pstats.Stats(self._prof).total_tt  # type: ignore
+        elif self._prof_type == "raw":
+            return self._prof.duration  # type: ignore