diff --git a/code/lib/preview-api/src/modules/addons/hooks.ts b/code/lib/preview-api/src/modules/addons/hooks.ts index 6c1c7847bbc8..c5f5acb3eec3 100644 --- a/code/lib/preview-api/src/modules/addons/hooks.ts +++ b/code/lib/preview-api/src/modules/addons/hooks.ts @@ -31,7 +31,7 @@ interface Effect { type AbstractFunction = (...args: any[]) => any; -export class HooksContext { +export class HooksContext { hookListsMap: WeakMap = undefined as any; mountedDecorators: Set = undefined as any; @@ -52,7 +52,7 @@ export class HooksContext { hasUpdates: boolean = undefined as any; - currentContext: StoryContext | null = undefined as any; + currentContext: StoryContext | null = undefined as any; renderListener = (storyId: StoryId) => { if (storyId !== this.currentContext?.id) { @@ -216,12 +216,18 @@ const areDepsEqual = (deps: any[], nextDeps: any[]) => const invalidHooksError = () => new Error('Storybook preview hooks can only be called inside decorators and story functions.'); -function getHooksContextOrNull(): HooksContext | null { +function getHooksContextOrNull< + TRenderer extends Renderer, + TArgs extends Args = Args +>(): HooksContext | null { return global.STORYBOOK_HOOKS_CONTEXT || null; } -function getHooksContextOrThrow(): HooksContext { - const hooks = getHooksContextOrNull(); +function getHooksContextOrThrow< + TRenderer extends Renderer, + TArgs extends Args = Args +>(): HooksContext { + const hooks = getHooksContextOrNull(); if (hooks == null) { throw invalidHooksError(); } @@ -405,8 +411,11 @@ export function useChannel(eventMap: EventMap, deps: any[] = []) { } /* Returns current story context */ -export function useStoryContext(): StoryContext { - const { currentContext } = getHooksContextOrThrow(); +export function useStoryContext< + TRenderer extends Renderer, + TArgs extends Args = Args +>(): StoryContext { + const { currentContext } = getHooksContextOrThrow(); if (currentContext == null) { throw invalidHooksError(); } @@ -424,26 +433,25 @@ export function useParameter(parameterKey: string, defaultValue?: S): S | und } /* Returns current value of story args */ -export function useArgs(): [ - SpecificArgs, - (newArgs: Partial) => void, - (argNames?: (keyof SpecificArgs)[]) => void +export function useArgs(): [ + TArgs, + (newArgs: Partial) => void, + (argNames?: (keyof TArgs)[]) => void ] { const channel = addons.getChannel(); - const { id: storyId, args } = useStoryContext(); + const { id: storyId, args } = useStoryContext(); const updateArgs = useCallback( - (updatedArgs: Partial) => - channel.emit(UPDATE_STORY_ARGS, { storyId, updatedArgs }), + (updatedArgs: Partial) => channel.emit(UPDATE_STORY_ARGS, { storyId, updatedArgs }), [channel, storyId] ); const resetArgs = useCallback( - (argNames?: (keyof SpecificArgs)[]) => channel.emit(RESET_STORY_ARGS, { storyId, argNames }), + (argNames?: (keyof TArgs)[]) => channel.emit(RESET_STORY_ARGS, { storyId, argNames }), [channel, storyId] ); - return [args as SpecificArgs, updateArgs, resetArgs]; + return [args as TArgs, updateArgs, resetArgs]; } /* Returns current value of global args */