fix timeit typing

uriva · Sep 19, 2023 · 462db0b · 462db0b
1 parent d946385
commit 462db0b
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 25 deletions.
diff --git a/src/debug.test.ts b/src/debug.test.ts
@@ -1,4 +1,4 @@
-import { assert, asyncTimeit, timeit } from "./debug.ts";
+import { assert, timeit } from "./debug.ts";
 import {
   assertSpyCall,
   assertSpyCalls,
@@ -21,7 +21,7 @@ Deno.test("timeit", () => {
 
 Deno.test("asyncTimeit", async () => {
   const logSpy = spy(console.log);
-  await asyncTimeit(
+  await timeit(
     (_, [time]) => logSpy(`slept for ${time}ms`),
     sleep,
   )(100);
@@ -35,7 +35,7 @@ Deno.test("asyncTimeit with args", async () => {
     await sleep(100);
     return a + b + c;
   };
-  await asyncTimeit(
+  await timeit(
     (_1, _2, result) => logSpy(`slept and returned ${result}`),
     f,
   )({ a: 1, b: 2, c: 3 });

diff --git a/src/debug.ts b/src/debug.ts
@@ -1,4 +1,5 @@
-import { Func } from "./typing.ts";
+import { Func, ReturnTypeUnwrapped } from "./typing.ts";
+
 import { sideEffect } from "./composition.ts";
 
 export const sideLog = <T>(x: T) => {
@@ -39,27 +40,28 @@ export const logWith = <T>(...x: any[]) =>
 
 const getTimestampMilliseconds = () => new Date().getTime();
 
-export const asyncTimeit = <Args extends unknown[], R>(
-  handler: (time: number, args: Args, result: R) => void,
-  f: (..._: Args) => R,
-) =>
-async (...args: Args) => {
-  const started = getTimestampMilliseconds();
-  const result = await f(...args);
-  handler(getTimestampMilliseconds() - started, args, result);
-  return result;
-};
-
-export const timeit = <Args extends unknown[], R>(
-  handler: (time: number, args: Args, result: R) => void,
-  f: (..._: Args) => R,
-) =>
-(...args: Args) => {
-  const started = getTimestampMilliseconds();
-  const result = f(...args);
-  handler(getTimestampMilliseconds() - started, args, result);
-  return result;
-};
+export const timeit = <F extends (..._: any[]) => any>(
+  handler: (
+    elapsed: number,
+    args: Parameters<F>,
+    result: ReturnTypeUnwrapped<F>,
+  ) => void,
+  f: F,
+): F =>
+  ((...x: Parameters<F>) => {
+    const started = getTimestampMilliseconds();
+    const result = f(...x);
+    if (result instanceof Promise) {
+      return result.then((result) => {
+        const elapsed = getTimestampMilliseconds() - started;
+        handler(elapsed, x, result);
+        return result;
+      });
+    }
+    const elapsed = getTimestampMilliseconds() - started;
+    handler(elapsed, x, result);
+    return result;
+  }) as F;
 
 export const assert = <T>(condition: (_: T) => boolean, errorMessage: string) =>
   sideEffect((x: T) => {