feat: statement / timeit nano benchmarks

caffeine/microbenchmarks/bench_redis.py

@@ -3,13 +3,9 @@
 import frappe
 
 from caffeine.microbenchmarks.bench_orm import get_all_roles
+from caffeine.microbenchmarks.utils import NanoBenchmark
 
-
-def bench_make_key():
-	keys = []
-	for dt in get_all_doctypes():
-		keys.append(frappe.cache.make_key(dt))
-	return keys
+bench_make_key = NanoBenchmark(statement="frappe.cache.make_key('a')")
 
 
 def bench_redis_get_set_delete_cycle():

caffeine/microbenchmarks/bench_utils.py

@@ -3,6 +3,9 @@
 import frappe
 from frappe.utils.caching import redis_cache, request_cache, site_cache
 
+# NOTE: decorated functions themselves are benchmarks, since they aren't wrapped by any other
+# function calls we don't need to move them to timeit style statements.
+
 
 @site_cache
 def bench_site_cache_no_arg():

caffeine/microbenchmarks/run_benchmarks.py

@@ -16,6 +16,7 @@
 	bench_utils,
 	bench_web_requests,
 )
+from caffeine.microbenchmarks.utils import NanoBenchmark
 
 BENCHMARK_PREFIX = "bench_"
 
@@ -40,8 +41,22 @@ def update_cmd_line(cmd, args):
 	args = runner.argparser.parse_args()
 	benchmarks = discover_benchmarks(cstr(args.benchmark_filter))
 	setup(args.site)
-	for name, func in benchmarks:
-		runner.bench_func(name, func)
+	default_globals = {"frappe": frappe}
+	for name, bench in benchmarks:
+		if isinstance(bench, FunctionType):
+			runner.bench_func(name, bench)
+		elif isinstance(bench, NanoBenchmark):
+			bench_globals = default_globals.copy()
+			bench_globals.update(bench.globals or {})
+			runner.timeit(
+				name,
+				stmt=bench.statement,
+				setup=bench.setup,
+				teardown=bench.teardown,
+				globals=bench_globals,
+			)
+		else:
+			raise ValueError("Unknown benchmark type:", type(bench))
 	teardown(args.site)
 
 
@@ -69,11 +84,13 @@ def discover_benchmarks(benchmark_filter):
 	benchmarks = []
 	for module in benchmark_modules:
 		module_name = module.__name__.split(".")[-1]
-		for fn_name, fn in inspect.getmembers(module, predicate=lambda x: isinstance(x, FunctionType)):
-			if fn_name.startswith(BENCHMARK_PREFIX):
-				unique_name = f"{module_name}_{fn.__name__}"
+		for bench_name, bench in inspect.getmembers(
+			module, predicate=lambda x: isinstance(x, FunctionType | NanoBenchmark)
+		):
+			if bench_name.startswith(BENCHMARK_PREFIX):
+				unique_name = f"{module_name}_{bench_name}"
 				if benchmark_filter in unique_name:
-					benchmarks.append((unique_name, fn))
+					benchmarks.append((unique_name, bench))
 
 	return sorted(benchmarks, key=lambda x: x[0])
 

caffeine/microbenchmarks/utils.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class NanoBenchmark:
+	statement: str
+	setup: str = "pass"
+	teardown: str = "pass"
+	globals: dict[str, Any] | None = None