From 6a16b9ce0f98babae42d245c3af6655e23eff8e2 Mon Sep 17 00:00:00 2001
From: xeho91 <xeho91@pm.me>
Date: Thu, 23 May 2024 17:13:00 +0800
Subject: [PATCH] Change `setTemplate` to use `typeof Story`

---
 src/components/Story.svelte         | 19 +++++++++----------
 src/components/StoryRenderer.svelte | 14 +++++++-------
 src/components/context.svelte.ts    | 16 ++++++++++------
 src/index.ts                        | 11 ++++++++---
 stories/Example.stories.svelte      |  6 +++---
 stories/templates.stories.svelte    |  8 ++++----
 6 files changed, 41 insertions(+), 33 deletions(-)

diff --git a/src/components/Story.svelte b/src/components/Story.svelte
index 6e83d08..662de31 100644
--- a/src/components/Story.svelte
+++ b/src/components/Story.svelte
@@ -1,12 +1,11 @@
 <script lang="ts" generics="M extends Meta">
-  import type { Meta, StoryObj } from '@storybook/svelte';
+  import type { Meta, StoryObj, StoryContext } from '@storybook/svelte';
   import type { Snippet } from 'svelte';
 
-  import type { TArgs, TContext } from '../index.js';
-  import {useStoriesExtractor, useStoryRenderer, useStoriesTemplate } from './context.svelte.js';
+  import { useStoriesExtractor, useStoryRenderer, useStoriesTemplate } from './context.svelte.js';
 
-  type Props = StoryObj<M> & {
-    children?: Snippet<[TArgs<M>, TContext<M>]>;
+  type Props = {
+    children?: Snippet<[StoryObj<M>["args"], StoryContext<M["args"]>]>;
     /**
     * Id of the story.
     *
@@ -33,13 +32,13 @@
      * If source is a string, it replaces the source of the story.
      */
     source?: boolean | string;
-  }
+  } & StoryObj<M>;
 
-  const { children, name = "Default",  id, play, ...restProps }: Props = $props();
+  const { children, name = "Default", id, play, ...restProps }: Props = $props();
 
   const extractor = useStoriesExtractor<M>();
   const renderer = useStoryRenderer<M>();
-  const storiesTemplate = useStoriesTemplate<M>();
+  const template = useStoriesTemplate<M>();
 
   const isCurrentlyViewed = $derived(!extractor.isExtracting && renderer.currentStoryName === name);
 
@@ -64,7 +63,7 @@
 {#if isCurrentlyViewed}
   {#if children}
     {@render children(renderer.args, renderer.storyContext)}
-  {:else if storiesTemplate}
-    {@render storiesTemplate(renderer.args, renderer.storyContext)}
+  {:else if template}
+    {@render template(renderer.args, renderer.storyContext)}
   {/if}
 {/if}
diff --git a/src/components/StoryRenderer.svelte b/src/components/StoryRenderer.svelte
index b1b2f5b..526e108 100644
--- a/src/components/StoryRenderer.svelte
+++ b/src/components/StoryRenderer.svelte
@@ -1,19 +1,19 @@
-<script lang="ts" generics="Component extends SvelteComponent">
-  import type { StoryContext } from '@storybook/svelte';
+<script lang="ts" generics="M extends Meta">
+  import type { Meta, StoryContext, StoryObj } from '@storybook/svelte';
   import type { StoryName } from '@storybook/types';
-  import type { ComponentProps, ComponentType, SvelteComponent } from 'svelte';
+  import type { ComponentType } from 'svelte';
 
   import { useStoryRenderer } from './context.svelte.js';
 
   type Props = {
-    Stories: Component extends SvelteComponent ? ComponentType<SvelteComponent> : never;
+    Stories: ComponentType;
     storyName: StoryName;
-    args: ComponentProps<Component>;
-    storyContext: StoryContext<ComponentProps<Component>>;
+    args: StoryObj<M>['args'];
+    storyContext: StoryContext<M['args']>;
   };
 
   let { Stories, storyName, args, storyContext }: Props = $props();
-  const context = useStoryRenderer<Component>();
+  const context = useStoryRenderer<M>();
 
   $effect(() => {
     context.set({
diff --git a/src/components/context.svelte.ts b/src/components/context.svelte.ts
index 66a87cf..ee3e9f7 100644
--- a/src/components/context.svelte.ts
+++ b/src/components/context.svelte.ts
@@ -1,8 +1,8 @@
-import type { Meta, StoryContext, StoryObj } from '@storybook/svelte';
+import type { Meta, StoryObj, StoryContext } from '@storybook/svelte';
 import type { StoryName } from '@storybook/types';
 import { getContext, hasContext, setContext, type Snippet } from 'svelte';
 
-import type { TArgs, TContext } from '../index.js';
+import type { Args, Story } from '../index.js';
 
 const KEYS = {
   extractor: 'storybook-stories-extractor-context',
@@ -121,7 +121,7 @@ export function useStoryRenderer<M extends Meta>() {
 }
 
 function createStoriesTemplateContext<M extends Meta>() {
-  let template = $state<Snippet<[TArgs<M>, TContext<M>]> | undefined>();
+  let template = $state<Snippet<[StoryObj<M>['args'], StoryContext<M['args']>]> | undefined>();
 
   function set(snippet?: typeof template) {
     template = snippet;
@@ -145,12 +145,16 @@ export function useStoriesTemplate<M extends Meta>() {
   return getContext<StoriesTemplateContext<M>>(KEYS.renderSnippet).template;
 }
 
-export function setTemplate<M extends Meta>(snippet?: StoriesTemplateContext<M>['template']): void {
+type InferMeta<S extends Story> = S extends Story<infer M extends Meta> ? M : never;
+
+export function setTemplate<S extends Story>(
+  snippet?: StoriesTemplateContext<InferMeta<S>>['template']
+): void {
   if (!hasContext(KEYS.renderSnippet)) {
-    setContext(KEYS.renderSnippet, createStoriesTemplateContext<M>());
+    setContext(KEYS.renderSnippet, createStoriesTemplateContext<InferMeta<S>>());
   }
 
-  const ctx = getContext<StoriesTemplateContext<M>>(KEYS.renderSnippet);
+  const ctx = getContext<StoriesTemplateContext<InferMeta<S>>>(KEYS.renderSnippet);
 
   ctx.set(snippet);
 }
diff --git a/src/index.ts b/src/index.ts
index a6d5fb0..d37e00e 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,9 +1,10 @@
 /// <reference types="webpack-env" />
 
-import type { Meta, StoryContext, StoryObj } from '@storybook/svelte';
+import type { Meta, StoryContext as _SB_StoryContext, StoryObj } from '@storybook/svelte';
 
 import Story from './components/Story.svelte';
 
+export type { Story };
 export { setTemplate } from './components/context.svelte.js';
 
 // FIXME: We don't use webpack anymore(?)
@@ -12,8 +13,12 @@ if (module?.hot?.decline) {
 }
 
 // TODO: Consult if these are correct types
-export type TArgs<M extends Meta> = StoryObj<M>['args'];
-export type TContext<M extends Meta> = StoryContext<M['args']>;
+export type Args<S extends Story> = S extends typeof Story<infer M extends Meta>
+  ? StoryObj<M>['args']
+  : never;
+export type StoryContext<S extends Story> = S extends typeof Story<infer M extends Meta>
+  ? _SB_StoryContext<M['args']>
+  : never;
 
 export function defineMeta<const M extends Meta>(meta: M) {
   return {
diff --git a/stories/Example.stories.svelte b/stories/Example.stories.svelte
index 0c99d72..50ed440 100644
--- a/stories/Example.stories.svelte
+++ b/stories/Example.stories.svelte
@@ -1,12 +1,12 @@
 <script context="module" lang="ts">
   import { action } from '@storybook/addon-actions';
 
-  import { defineMeta, setTemplate, type TArgs, type TContext } from "../src/index.js";
+  import { defineMeta, setTemplate, type Args, type StoryContext } from "../src/index.js";
 
   import Example from './Example.svelte';
 
   // Description set explicitly in the comment above `defineMeta`
-  const { Story, meta: m } = defineMeta({
+  const { Story } = defineMeta({
     title: 'Example',
     component: Example,
     tags: ['autodocs'],
@@ -29,7 +29,7 @@
   setTemplate(render);
 </script>
 
-{#snippet render(args: TArgs<typeof m>, context: TContext<typeof m>)}
+{#snippet render(args: Args<typeof Story>, context: StoryContext<typeof Story>)}
     <Example {...args} onclick={handleClick}>
       <p>{context.name}</p>
       You clicked: {count}<br>
diff --git a/stories/templates.stories.svelte b/stories/templates.stories.svelte
index f03ecc9..c858661 100644
--- a/stories/templates.stories.svelte
+++ b/stories/templates.stories.svelte
@@ -1,5 +1,5 @@
 <script context="module" lang="ts">
-  import { defineMeta, type TArgs } from '../src/index.js';
+  import { defineMeta, type Args } from '../src/index.js';
 
   import Text from './Text.svelte';
 
@@ -7,7 +7,7 @@
    * Demonstration on how to use multiple templates in one stories file,
    * powered by Svelte's **snippets**.
    */
-  const { Story, meta } = defineMeta({
+  const { Story } = defineMeta({
     title: 'Templates',
     component: Text,
     tags: ['autodocs'],
@@ -18,12 +18,12 @@
   });
 </script>
 
-{#snippet template1(args: TArgs<typeof meta>)}
+{#snippet template1(args: Args<typeof Story>)}
   <h2 style="color: lightgreen">Template 1</h2>
   <p>{args?.text}</p>
 {/snippet}
 
-{#snippet template2(args: TArgs<typeof meta>)}
+{#snippet template2(args: Args<typeof Story>)}
   <h2 style="color: fuchsia">Template 2</h2>
   <hr>
   <p>{args?.text}</p>