diff --git a/packages/core/src/except.ts b/packages/core/src/except.ts new file mode 100644 index 00000000000..7472721f85e --- /dev/null +++ b/packages/core/src/except.ts @@ -0,0 +1,13 @@ +// Do not convert this file to .d.ts + +type IsEqual = + (() => G extends T ? 1 : 2) extends + (() => G extends U ? 1 : 2) + ? true + : false; + +type Filter = IsEqual extends true ? never : (KeyType extends ExcludeType ? never : KeyType); + +export type Except = { + [KeyType in keyof ObjectType as Filter]: ObjectType[KeyType]; +}; diff --git a/packages/core/src/router.ts b/packages/core/src/router.ts index 4ba52e890d5..96b4661e26c 100644 --- a/packages/core/src/router.ts +++ b/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; } +type GetUserDataFromRequest = T extends Request ? Y : never; + /** * Simple router that works based on request labels. This instance can then serve as a `requestHandler` of your crawler. * @@ -72,7 +77,10 @@ export class Router { /** * Registers new route handler for given label. */ - addHandler(label: string | symbol, handler: (ctx: Context) => Awaitable) { + addHandler>( + label: string | symbol, + handler: (ctx: Except & {request: Request}) => Awaitable, + ) { this.validate(label); this.routes.set(label, handler); } @@ -80,7 +88,9 @@ export class Router { /** * Registers default route handler. */ - addDefaultHandler(handler: (ctx: Context) => Awaitable) { + addDefaultHandler>( + handler: (ctx: Except & {request: Request}) => Awaitable, + ) { this.validate(defaultRoute); this.routes.set(defaultRoute, handler); } diff --git a/test/core/router.test.ts b/test/core/router.test.ts index 026a4751ff5..2ba884585e2 100644 --- a/test/core/router.test.ts +++ b/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): void => {}; + + const router: Router> = { + 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); + }); + }); });