🚧 Add extract-instance-nodes & extract snippet for setTemplate

src/parser/extract/svelte/fragment-nodes.ts

@@ -1,18 +1,21 @@
-import type { Comment, Component, Fragment, SvelteNode } from 'svelte/compiler';
+import type { Comment, Component, Fragment, SnippetBlock, SvelteNode } from 'svelte/compiler';
 import type { Visitors } from 'zimmerframe';
 
 import type { extractModuleNodes } from './module-nodes.js';
+import type { extractInstanceNodes } from './instance-nodes.js';
 
 interface Result {
   storyComponents: Array<{
     comment?: Comment;
     component: Component;
   }>;
+  setTemplateSnippet: SnippetBlock | undefined;
 }
 
 interface Params {
   fragment: Fragment;
   filename?: string;
+  instanceNodes: Awaited<ReturnType<typeof extractInstanceNodes>>;
   moduleNodes: Awaited<ReturnType<typeof extractModuleNodes>>;
 }
 
@@ -24,13 +27,15 @@ interface Params {
 export async function extractFragmentNodes(params: Params): Promise<Result> {
   const { walk } = await import('zimmerframe');
 
-  const { fragment, filename, moduleNodes } = params;
+  const { fragment, filename, moduleNodes, instanceNodes } = params;
+  const { setTemplateCall } = instanceNodes;
   const { storyIdentifier } = moduleNodes;
 
   let latestComment: Comment | undefined;
 
   const state: Result = {
     storyComponents: [],
+    setTemplateSnippet: undefined,
   };
 
   const visitors: Visitors<SvelteNode, typeof state> = {
@@ -48,6 +53,16 @@ export async function extractFragmentNodes(params: Params): Promise<Result> {
         latestComment = undefined;
       }
     },
+
+    SnippetBlock(node, { state }) {
+      if (
+        setTemplateCall &&
+        setTemplateCall.arguments[0].type === 'Identifier' &&
+        setTemplateCall.arguments[0].name === node.expression.name
+      ) {
+        state.setTemplateSnippet = node;
+      }
+    },
   };
 
   walk(fragment, state, visitors);

src/parser/extract/svelte/instance-nodes.ts

@@ -0,0 +1,50 @@
+import type { CallExpression } from 'estree';
+import type { Root, SvelteNode } from 'svelte/compiler';
+import type { Visitors } from 'zimmerframe';
+
+import type { extractModuleNodes } from './module-nodes.js';
+
+interface Result {
+  setTemplateCall: CallExpression | undefined;
+}
+
+interface Params {
+  instance: Root['instance'];
+  moduleNodes: Awaited<ReturnType<typeof extractModuleNodes>>;
+  filename?: string;
+}
+
+/**
+ * Extract Svelte AST nodes via `svelte.compile`,
+ * and from the instance tag - `<script>` _(without `context="module"`)_.
+ * They are needed for further code analysis/transformation.
+	// NOTE: Is optional for the `*.stories.svelte` files to have this tag.
+ */
+export async function extractInstanceNodes(options: Params): Promise<Result> {
+  const { instance, moduleNodes, filename } = options;
+  const { setTemplateImport } = moduleNodes;
+
+  if (!instance || !setTemplateImport) {
+    return {
+      setTemplateCall: undefined,
+    };
+  }
+  const { walk } = await import('zimmerframe');
+
+  const state: Partial<Result> = {};
+  const visitors: Visitors<SvelteNode, typeof state> = {
+    CallExpression(node, { state }) {
+      if (node.callee.type === 'Identifier' && node.callee.name === setTemplateImport?.local.name) {
+        state.setTemplateCall = node;
+      }
+    },
+  };
+
+  walk(instance.content, state, visitors);
+
+  const { setTemplateCall } = state;
+
+  return {
+    setTemplateCall,
+  };
+}

src/parser/extract/svelte/nodes.ts

@@ -2,13 +2,16 @@ import type { Root } from 'svelte/compiler';
 
 import { extractModuleNodes } from './module-nodes.js';
 import { extractFragmentNodes } from './fragment-nodes.js';
+import { extractInstanceNodes } from './instance-nodes.js';
 
 /**
  * Selected nodes extracted from the Svelte AST via `svelte.compile`,
  * needed for further code analysis/transformation.
  */
 export type SvelteASTNodes = Awaited<ReturnType<typeof extractModuleNodes>> &
-  Awaited<ReturnType<typeof extractFragmentNodes>>;
+  Awaited<
+    ReturnType<typeof extractFragmentNodes> & Awaited<ReturnType<typeof extractInstanceNodes>>
+  >;
 
 interface Params {
   ast: Root;
@@ -20,17 +23,24 @@ interface Params {
  */
 export async function extractSvelteASTNodes(params: Params): Promise<SvelteASTNodes> {
   const { ast, filename } = params;
-  const { module, fragment } = ast;
+  const { module, fragment, instance } = ast;
 
   const moduleNodes = await extractModuleNodes({ module, filename });
+  const instanceNodes = await extractInstanceNodes({
+    instance,
+    filename,
+    moduleNodes,
+  });
   const fragmentNodes = await extractFragmentNodes({
     fragment,
     filename,
     moduleNodes,
+    instanceNodes,
   });
 
   return {
     ...moduleNodes,
+    ...instanceNodes,
     ...fragmentNodes,
   };
 }