feat: make piscina entrypoint easier for typescript

Author: danielpza

docs/docs/getting-started/typescript.mdx

@@ -112,3 +112,32 @@ const piscina = new Piscina({
 })();
 ```
 
+## Using generics
+
+You can use generics to the types for the worker function:
+
+```typescript title='worker.ts'
+
+interface Inputs {
+  a: number;
+  b: number;
+}
+
+export default function addNumbers({ a, b }: Inputs): number {
+  return a + b;
+}
+```
+
+```typescript title='main.ts'
+
+import Piscina from 'piscina';
+
+const piscina = new Piscina<typeof import("./worker.ts")>({
+  filename: new URL("./worker.ts", import.meta.url).href,
+});
+
+(async () => {
+  const result = await piscina.run({ a: 2, b: 3 });
+  console.log('Result:', result);
+})();
+```

package-lock.json

@@ -51,6 +51,7 @@
     "c8": "^10.1.2",
     "concat-stream": "^2.0.0",
     "eslint": "^9.16.0",
+    "expect-type": "^1.2.1",
     "gen-esm-wrapper": "^1.1.1",
     "neostandard": "^0.12.0",
     "tap": "^16.3.7",

package.json

@@ -97,18 +97,18 @@ interface FilledOptions extends Options {
   workerHistogram: boolean,
 }
 
-interface RunOptions {
+interface RunOptions<Func = string> {
   transferList? : TransferList,
   filename? : string | null,
   signal? : AbortSignalAny | null,
-  name? : string | null
+  name? : Func | null
 }
 
-interface FilledRunOptions extends RunOptions {
+interface FilledRunOptions<Func extends string = string> extends RunOptions<Func> {
   transferList : TransferList | never,
   filename : string | null,
   signal : AbortSignalAny | null,
-  name : string | null
+  name : Func | null
 }
 
 interface CloseOptions {
@@ -481,7 +481,7 @@ class ThreadPool {
 
   runTask (
     task : any,
-    options : RunOptions) : Promise<any> {
+    options : RunOptions<string>) : Promise<any> {
     let {
       filename,
       name
@@ -720,7 +720,7 @@ class ThreadPool {
   }
 }
 
-export default class Piscina<T = any, R = any> extends EventEmitterAsyncResource {
+export default class Piscina<Exports extends Record<string, (payload: any) => any> = Record<string, (payload: any) => any>> extends EventEmitterAsyncResource {
   #pool : ThreadPool;
   #histogram: PiscinaHistogram | null = null;
 
@@ -793,7 +793,7 @@ export default class Piscina<T = any, R = any> extends EventEmitterAsyncResource
     this.#pool = new ThreadPool(this, options);
   }
 
-  run (task : T, options : RunOptions = kDefaultRunOptions): Promise<R> {
+  run <Func extends keyof Exports = 'default'>(task : Parameters<Exports[Func]>[0], options : RunOptions<Func> = kDefaultRunOptions as RunOptions<Func>): Promise<ReturnType<Exports[Func]>> {
     if (options === null || typeof options !== 'object') {
       return Promise.reject(
         new TypeError('options must be an object'));

src/index.ts

@@ -0,0 +1,7 @@
+export function greet({ name }: { name: string }): string {
+  return `Hello ${name}!`;
+}
+
+export default function add({ a, b }: { a: number; b: number }): number {
+  return a + b;
+}

test/fixtures/worker-type.ts

@@ -0,0 +1,31 @@
+import { resolve } from "node:path";
+import { expectTypeOf } from "expect-type";
+
+import { test } from "tap";
+
+import Piscina from "../dist";
+
+test("types can be inferred from function", async (t) => {
+  const pool = new Piscina<typeof import("./fixtures/worker-type.js")>({
+    filename: resolve(__dirname, "fixtures/worker-type.js"),
+    concurrentTasksPerWorker: 1,
+  });
+
+  expectTypeOf(pool.run).toExtend<(payload: { a: number; b: number }) => Promise<number>>();
+  expectTypeOf(pool.run).toExtend<(payload: { a: number; b: number }, opts: { name: "default" }) => Promise<number>>();
+  expectTypeOf(pool.run).toExtend<(payload: { name: string }, opts: { name: "greet" }) => Promise<string>>();
+});
+
+test("types can be manually specified", async (t) => {
+  const pool = new Piscina<{
+    default: (payload: { a: number; b: number }) => number;
+    greet: (payload: { name: string }) => string;
+  }>({
+    filename: resolve(__dirname, "fixtures/worker-type.js"),
+    concurrentTasksPerWorker: 1,
+  });
+
+  expectTypeOf(pool.run).toExtend<(payload: { a: number; b: number }) => Promise<number>>();
+  expectTypeOf(pool.run).toExtend<(payload: { a: number; b: number }, opts: { name: "default" }) => Promise<number>>();
+  expectTypeOf(pool.run).toExtend<(payload: { name: string }, opts: { name: "greet" }) => Promise<string>>();
+});