Addon-docs: Reuse extractSource from source-loader

addons/docs/src/blocks/enhanceSource.ts

@@ -1,14 +1,10 @@
 import { combineParameters } from '@storybook/client-api';
 import { StoryContext, Parameters } from '@storybook/addons';
-
-interface Location {
-  line: number;
-  col: number;
-}
+import { extractSource, LocationsMap } from '@storybook/source-loader';
 
 interface StorySource {
   source: string;
-  locationsMap: { [id: string]: { startBody: Location; endBody: Location } };
+  locationsMap: LocationsMap;
 }
 
 /**
@@ -24,23 +20,9 @@ const extract = (targetId: string, { source, locationsMap }: StorySource) => {
 
   const sanitizedStoryName = storyIdToSanitizedStoryName(targetId);
   const location = locationsMap[sanitizedStoryName];
-
-  const { startBody: start, endBody: end } = location;
   const lines = source.split('\n');
-  if (start.line === end.line && lines[start.line - 1] !== undefined) {
-    return lines[start.line - 1].substring(start.col, end.col);
-  }
-  // NOTE: storysource locations are 1-based not 0-based!
-  const startLine = lines[start.line - 1];
-  const endLine = lines[end.line - 1];
-  if (startLine === undefined || endLine === undefined) {
-    return source;
-  }
-  return [
-    startLine.substring(start.col),
-    ...lines.slice(start.line, end.line - 1),
-    endLine.substring(0, end.col),
-  ].join('\n');
+
+  return extractSource(location, lines);
 };
 
 export const enhanceSource = (context: StoryContext): Parameters => {

lib/source-loader/src/abstract-syntax-tree/generate-helpers.js

@@ -9,6 +9,7 @@ import {
   popParametersObjectFromDefaultExport,
   findExportsMap as generateExportsMap,
 } from './traverse-helpers';
+import { extractSource } from '../extract-source';
 
 export function sanitizeSource(source) {
   return JSON.stringify(source)
@@ -176,27 +177,6 @@ export function generateSourcesInExportedParameters(source, ast, additionalParam
   return source;
 }
 
-/**
- * given a location, extract the text from the full source
- */
-function extractSource(location, lines) {
-  const { startBody: start, endBody: end } = location;
-  if (start.line === end.line && lines[start.line - 1] !== undefined) {
-    return lines[start.line - 1].substring(start.col, end.col);
-  }
-  // NOTE: storysource locations are 1-based not 0-based!
-  const startLine = lines[start.line - 1];
-  const endLine = lines[end.line - 1];
-  if (startLine === undefined || endLine === undefined) {
-    return null;
-  }
-  return [
-    startLine.substring(start.col),
-    ...lines.slice(start.line, end.line - 1),
-    endLine.substring(0, end.col),
-  ].join('\n');
-}
-
 function addStorySourceParameter(key, snippet) {
   const source = sanitizeSource(snippet);
   return `${key}.parameters = { storySource: { source: ${source} }, ...${key}.parameters };`;

lib/source-loader/src/extract-source.ts

@@ -0,0 +1,22 @@
+import type { SourceBlock } from './types';
+
+/**
+ * given a location, extract the text from the full source
+ */
+export function extractSource(location: SourceBlock, lines: string[]): string | null {
+  const { startBody: start, endBody: end } = location;
+  if (start.line === end.line && lines[start.line - 1] !== undefined) {
+    return lines[start.line - 1].substring(start.col, end.col);
+  }
+  // NOTE: storysource locations are 1-based not 0-based!
+  const startLine = lines[start.line - 1];
+  const endLine = lines[end.line - 1];
+  if (startLine === undefined || endLine === undefined) {
+    return null;
+  }
+  return [
+    startLine.substring(start.col),
+    ...lines.slice(start.line, end.line - 1),
+    endLine.substring(0, end.col),
+  ].join('\n');
+}

lib/source-loader/src/index.ts

@@ -4,3 +4,5 @@ import { transform } from './build';
 export default transform;
 
 export * from './types';
+
+export { extractSource } from './extract-source';

lib/source-loader/src/types.ts

@@ -4,6 +4,8 @@ export interface SourceLoc {
 }
 
 export interface SourceBlock {
+  startBody: SourceLoc;
+  endBody: SourceLoc;
   startLoc: SourceLoc;
   endLoc: SourceLoc;
 }