feat(router): add userdata generic to addHandler (#1547)

packages/core/src/except.ts

@@ -0,0 +1,13 @@
+// Do not convert this file to .d.ts
+
+type IsEqual<T, U> =
+    (<G>() => G extends T ? 1 : 2) extends
+    (<G>() => G extends U ? 1 : 2)
+        ? true
+        : false;
+
+type Filter<KeyType, ExcludeType> = IsEqual<KeyType, ExcludeType> extends true ? never : (KeyType extends ExcludeType ? never : KeyType);
+
+export type Except<ObjectType, KeysType extends keyof ObjectType> = {
+    [KeyType in keyof ObjectType as Filter<KeyType, KeysType>]: ObjectType[KeyType];
+};

packages/core/src/router.ts

@@ -1,5 +1,8 @@
+import type { Dictionary } from '@crawlee/types';
 import type { CrawlingContext } from './crawlers/crawler_commons';
 import type { Awaitable } from './typedefs';
+import type { Request } from './request';
+import type { Except } from './except';
 import { MissingRouteError } from './errors';
 
 const defaultRoute = Symbol('default-route');
@@ -8,6 +11,8 @@ export interface RouterHandler<Context extends CrawlingContext = CrawlingContext
     (ctx: Context): Awaitable<void>;
 }
 
+type GetUserDataFromRequest<T> = T extends Request<infer Y> ? Y : never;
+
 /**
  * Simple router that works based on request labels. This instance can then serve as a `requestHandler` of your crawler.
  *
@@ -72,15 +77,20 @@ export class Router<Context extends CrawlingContext> {
     /**
      * Registers new route handler for given label.
      */
-    addHandler(label: string | symbol, handler: (ctx: Context) => Awaitable<void>) {
+    addHandler<UserData extends Dictionary = GetUserDataFromRequest<Context['request']>>(
+        label: string | symbol,
+        handler: (ctx: Except<Context, 'request'> & {request: Request<UserData>}) => Awaitable<void>,
+    ) {
         this.validate(label);
         this.routes.set(label, handler);
     }
 
     /**
      * Registers default route handler.
      */
-    addDefaultHandler(handler: (ctx: Context) => Awaitable<void>) {
+    addDefaultHandler<UserData extends Dictionary = GetUserDataFromRequest<Context['request']>>(
+        handler: (ctx: Except<Context, 'request'> & {request: Request<UserData>}) => Awaitable<void>,
+    ) {
         this.validate(defaultRoute);
         this.routes.set(defaultRoute, handler);
     }

test/core/router.test.ts

@@ -1,3 +1,4 @@
+import type { CrawlingContext } from '@crawlee/core';
 import { MissingRouteError, Router } from '@crawlee/core';
 import { BasicCrawler } from '@crawlee/basic';
 
@@ -104,4 +105,29 @@ describe('Router', () => {
             requestHandler: router,
         });
     });
+
+    test('addHandler accepts userdata generic', async () => {
+        const testType = <T>(t: T): void => {};
+
+        const router: Router<CrawlingContext<{foo: 'foo'}>> = {
+            addHandler: () => {},
+            addDefaultHandler: () => {},
+        } as any;
+
+        router.addHandler('1', (ctx) => {
+            testType<'foo'>(ctx.request.userData.foo);
+        });
+
+        router.addHandler<{foo: 'bar'}>('2', (ctx) => {
+            testType<'bar'>(ctx.request.userData.foo);
+        });
+
+        router.addDefaultHandler((ctx) => {
+            testType<'foo'>(ctx.request.userData.foo);
+        });
+
+        router.addDefaultHandler<{foo: 'bar'}>((ctx) => {
+            testType<'bar'>(ctx.request.userData.foo);
+        });
+    });
 });