Add story source code to components.json (#2948)

* Add story source code to component.json

* Create .changeset/dry-cows-try.md

* Bump output schema version

* Update generated/components.json

* Remove semi from story code

* Update generated/components.json

* Handle invalid story ids

* Trim leading semi

* Update schema

* Update generated/components.json

* Add more TreeView stories

* Update generated/components.json

---------

Co-authored-by: colebemis <colebemis@users.noreply.github.com>

.changeset/dry-cows-try.md

@@ -0,0 +1,5 @@
+---
+"@primer/react": patch
+---
+
+Include story source code in `generated/components.json`

script/components-json/build.ts

@@ -1,31 +1,88 @@
-import fs from 'fs'
+import generate from '@babel/generator'
+import {parse} from '@babel/parser'
+import traverse from '@babel/traverse'
+import {ArrowFunctionExpression, Identifier, VariableDeclaration} from '@babel/types'
+import Ajv from 'ajv'
+import {pascalCase} from 'change-case'
 import glob from 'fast-glob'
+import fs from 'fs'
 import keyBy from 'lodash.keyby'
-import Ajv from 'ajv'
+import prettier from 'prettier'
 import componentSchema from './component.schema.json'
 import outputSchema from './output.schema.json'
 
+// Only includes fields we use in this script
+type Component = {
+  name: string
+  stories: Array<{id: string; code?: string}>
+}
+
 const ajv = new Ajv()
 
-const jsonFiles = glob.sync('src/**/*.docs.json')
+// Get all JSON files matching `src/**/*.docs.json`
+const docsFiles = glob.sync('src/**/*.docs.json')
+
+const components = docsFiles.map(docsFilepath => {
+  const docs = JSON.parse(fs.readFileSync(docsFilepath, 'utf-8'))
+
+  // Create a validator for the component schema
+  const validate = ajv.compile<Component>(componentSchema)
+
+  // Validate the component schema
+  if (!validate(docs)) {
+    throw new Error(`Invalid docs file ${docsFilepath}: ${JSON.stringify(validate.errors, null, 2)}}`)
+  }
+
+  // Get path to default story file
+  // Example: src/components/Box/Box.docs.json -> src/components/Box/Box.stories.tsx
+  const defaultStoryFilepath = docsFilepath.replace(/\.docs\.json$/, '.stories.tsx')
+
+  // Get the default story id
+  const defaultStoryId = `components-${String(docs.name).toLowerCase()}--default`
+
+  // Get source code for default story
+  const {Default: defaultStoryCode} = getStorySourceCode(defaultStoryFilepath)
+
+  // Get path to feature story file
+  // Example: src/components/Box/Box.docs.json -> src/components/Box/Box.features.stories.tsx
+  const featureStoryFilepath = docsFilepath.replace(/\.docs\.json$/, '.features.stories.tsx')
+
+  // Get source code for each feature story
+  const featureStorySourceCode = getStorySourceCode(featureStoryFilepath)
 
-const components = jsonFiles.map(file => {
-  const json = JSON.parse(fs.readFileSync(file, 'utf-8'))
+  // Populate source code for each feature story
+  const featureStories = docs.stories
+    // Filter out the default story
+    .filter(({id}) => id !== defaultStoryId)
+    .map(({id}) => {
+      const storyName = getStoryName(id)
+      const code = featureStorySourceCode[storyName]
 
-  // Validate component json
-  const validate = ajv.compile(componentSchema)
+      if (!code) {
+        throw new Error(
+          `Invalid story id "${id}" in ${docsFilepath}. No story named "${storyName}" found in ${featureStoryFilepath}`,
+        )
+      }
 
-  if (!validate(json)) {
-    throw new Error(`Invalid file ${file}: ${JSON.stringify(validate.errors, null, 2)}}`)
+      return {id, code}
+    })
+
+  // Replace the stories array with the new array that includes source code
+  docs.stories = featureStories
+
+  // Add default story to the beginning of the array
+  if (defaultStoryCode) {
+    docs.stories.unshift({
+      id: defaultStoryId,
+      code: defaultStoryCode,
+    })
   }
 
-  // TODO: Assert that component ids use kebab-case
   // TODO: Provide default type and description for sx and ref props
-  // TODO: Sort component props
-  return json
+  return docs
 })
 
-const data = {schemaVersion: 1, components: keyBy(components, 'id')}
+const data = {schemaVersion: 2, components: keyBy(components, 'id')}
 
 // Validate output
 const validate = ajv.compile(outputSchema)
@@ -45,5 +102,60 @@ fs.writeFileSync('generated/components.json', JSON.stringify(data, null, 2))
 // Print success message
 // eslint-disable-next-line no-console
 console.log(
-  `Successfully built "generated/components.json" from ${jsonFiles.length} files matching "src/**/*.docs.json"`,
+  `Successfully built "generated/components.json" from ${docsFiles.length} files matching "src/**/*.docs.json"`,
 )
+
+// -----------------------------------------------------------------------------
+// Helper functions
+// -----------------------------------------------------------------------------
+
+/**
+ * Returns an object mapping story names to their source code
+ *
+ * @example
+ * getStorySourceCode('src/components/Button/Button.stories.tsx')
+ * // {Default: '<Button>Button</Button>'}
+ */
+function getStorySourceCode(filepath: string) {
+  if (!fs.existsSync(filepath)) {
+    return {}
+  }
+
+  const code = fs.readFileSync(filepath, 'utf-8')
+
+  const ast = parse(code, {
+    sourceType: 'module',
+    plugins: ['jsx', 'typescript'],
+  })
+
+  const stories: Record<string, string> = {}
+
+  traverse(ast, {
+    ExportNamedDeclaration(path) {
+      const varDecloration = path.node.declaration as VariableDeclaration
+      const id = varDecloration.declarations[0].id as Identifier
+      const name = id.name
+      const func = varDecloration.declarations[0].init as ArrowFunctionExpression
+
+      const code = prettier
+        .format(generate(func).code, {
+          parser: 'typescript',
+          singleQuote: true,
+          trailingComma: 'all',
+          semi: false,
+        })
+        .trim()
+        .replace(/;$/, '')
+        .replace(/^;/, '')
+
+      stories[name] = code
+    },
+  })
+
+  return stories
+}
+
+function getStoryName(id: string) {
+  const parts = id.split('--')
+  return pascalCase(parts[parts.length - 1])
+}

script/components-json/component.schema.json

@@ -73,7 +73,19 @@
       "type": "array",
       "description": "An array of Storybook story IDs to embed in the docs.",
       "items": {
-        "type": "string"
+        "type": "object",
+        "required": ["id"],
+        "additionalProperties": false,
+        "properties": {
+          "id": {
+            "type": "string",
+            "description": "The Storybook story ID (e.g. \"components-treeview--default\")."
+          },
+          "code": {
+            "type": "string",
+            "description": "The source code of the story. This will be automatically populated during build time."
+          }
+        }
       }
     },
     "props": {

script/components-json/output.schema.json

@@ -5,7 +5,7 @@
   "properties": {
     "schemaVersion": {
       "type": "number",
-      "enum": [1],
+      "enum": [2],
       "description": "The version of the schema. We increment this when we make breaking changes to the schema."
     },
     "components": {

src/TreeView/TreeView.docs.json

@@ -3,7 +3,14 @@
   "name": "TreeView",
   "status": "beta",
   "a11yReviewed": true,
-  "stories": ["components-treeview-features--files", "components-treeview-features--files-changed"],
+  "stories": [
+    {"id": "components-treeview-features--files"},
+    {"id": "components-treeview-features--files-changed"},
+    {"id": "components-treeview-features--async-success"},
+    {"id": "components-treeview-features--async-error"},
+    {"id": "components-treeview-features--async-with-count"},
+    {"id": "components-treeview-features--controlled"}
+  ],
   "props": [
     {
       "name": "children",