From cdb507992d600137678fd16924489246ddfb6b70 Mon Sep 17 00:00:00 2001
From: Valentin Palkovic <valentin@chromatic.com>
Date: Mon, 4 Dec 2023 09:04:41 +0100
Subject: [PATCH] Fix CSF Plugin for Angular

---
 code/lib/csf-plugin/src/index.ts         |  12 +-
 code/lib/csf-tools/src/enrichCsf.test.ts | 392 +++++++++++++++++------
 code/lib/csf-tools/src/enrichCsf.ts      |  22 +-
 3 files changed, 318 insertions(+), 108 deletions(-)

diff --git a/code/lib/csf-plugin/src/index.ts b/code/lib/csf-plugin/src/index.ts
index aed7a531ee32..962a192bd746 100644
--- a/code/lib/csf-plugin/src/index.ts
+++ b/code/lib/csf-plugin/src/index.ts
@@ -12,15 +12,17 @@ const logger = console;
 export const unplugin = createUnplugin<CsfPluginOptions>((options) => {
   return {
     name: 'unplugin-csf',
-    enforce: 'pre',
-    loadInclude(id) {
+    transformInclude(id) {
       return STORIES_REGEX.test(id);
     },
-    async load(fname) {
-      const code = await fs.readFile(fname, 'utf-8');
+    async transform(code, id) {
+      const sourceCode = await fs.readFile(id, 'utf-8');
       try {
         const csf = loadCsf(code, { makeTitle: (userTitle) => userTitle || 'default' }).parse();
-        enrichCsf(csf, options);
+        const csfSource = loadCsf(sourceCode, {
+          makeTitle: (userTitle) => userTitle || 'default',
+        }).parse();
+        enrichCsf(csf, csfSource, options);
         return formatCsf(csf, { sourceMaps: true });
       } catch (err: any) {
         // This can be called on legacy storiesOf files, so just ignore
diff --git a/code/lib/csf-tools/src/enrichCsf.test.ts b/code/lib/csf-tools/src/enrichCsf.test.ts
index a8f5f3aaa09c..d793f730cecc 100644
--- a/code/lib/csf-tools/src/enrichCsf.test.ts
+++ b/code/lib/csf-tools/src/enrichCsf.test.ts
@@ -11,11 +11,16 @@ expect.addSnapshotSerializer({
   test: (val) => true,
 });
 
-const enrich = (code: string, options?: EnrichCsfOptions) => {
+const enrich = (code: string, originalCode: string, options?: EnrichCsfOptions) => {
   // we don't actually care about the title
 
-  const csf = loadCsf(code, { makeTitle: (userTitle) => userTitle || 'default' }).parse();
-  enrichCsf(csf, options);
+  const csf = loadCsf(code, {
+    makeTitle: (userTitle) => userTitle || 'default',
+  }).parse();
+  const csfSource = loadCsf(originalCode, {
+    makeTitle: (userTitle) => userTitle || 'default',
+  }).parse();
+  enrichCsf(csf, csfSource, options);
   return formatCsf(csf);
 };
 
@@ -23,17 +28,28 @@ describe('enrichCsf', () => {
   describe('source', () => {
     it('csf1', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+          // compiled code
           export default {
            title: 'Button',
           }
-          export const Basic = () => <Button />
-        `)
+          export const Basic = () => React.createElement(Button)
+        `,
+          dedent`
+        // original code
+        export default {
+         title: 'Button',
+        }
+        export const Basic = () => <Button />
+      `
+        )
       ).toMatchInlineSnapshot(`
+        // compiled code
         export default {
           title: 'Button'
         };
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -48,19 +64,32 @@ describe('enrichCsf', () => {
     });
     it('csf2', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+        // compiled code
+        export default {
+          title: 'Button',
+        }
+        const Template = (args) => React.createElement(Button, args);
+        export const Basic = Template.bind({});
+        Basic.parameters = { foo: 'bar' }
+      `,
+          dedent`
+          // original code
           export default {
             title: 'Button',
           }
           const Template = (args) => <Button {...args} />
           export const Basic = Template.bind({});
           Basic.parameters = { foo: 'bar' }
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
+        // compiled code
         export default {
           title: 'Button'
         };
-        const Template = args => <Button {...args} />;
+        const Template = args => React.createElement(Button, args);
         export const Basic = Template.bind({});
         Basic.parameters = {
           foo: 'bar'
@@ -79,15 +108,26 @@ describe('enrichCsf', () => {
     });
     it('csf3', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+        // compiled code
+        export default {
+          title: 'Button',
+        }
+        export const Basic = { parameters: { foo: 'bar' } }
+      `,
+          dedent`
+          // original code
           export default {
             title: 'Button',
           }
           export const Basic = {
             parameters: { foo: 'bar' }
           }
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
+        // compiled code
         export default {
           title: 'Button'
         };
@@ -110,14 +150,26 @@ describe('enrichCsf', () => {
     });
     it('multiple stories', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+        // compiled code
+        export default {
+          title: 'Button',
+        }
+        export const A = {}
+        export const B = {}
+      `,
+          dedent`
+          // original code
           export default {
             title: 'Button',
           }
           export const A = {}
           export const B = {}
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
+        // compiled code
         export default {
           title: 'Button'
         };
@@ -150,19 +202,31 @@ describe('enrichCsf', () => {
   describe('story descriptions', () => {
     it('skips inline comments', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+        // compiled code
+        export default {
+         title: 'Button',
+        }
+        // The most basic button
+        export const Basic = () => React.createElement(Button);
+      `,
+          dedent`
+          // original code
           export default {
            title: 'Button',
           }
           // The most basic button
           export const Basic = () => <Button />
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
+        // compiled code
         export default {
           title: 'Button'
         };
         // The most basic button
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -178,19 +242,29 @@ describe('enrichCsf', () => {
 
     it('skips blocks without jsdoc', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+          // compiled code
+          export default {
+           title: 'Button',
+          }
+          export const Basic = () => React.createElement(Button)
+        `,
+          dedent`
+          // original code
           export default {
            title: 'Button',
           }
           /* The most basic button */
           export const Basic = () => <Button />
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
+        // compiled code
         export default {
           title: 'Button'
         };
-        /* The most basic button */
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -206,19 +280,29 @@ describe('enrichCsf', () => {
 
     it('JSDoc single-line', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+          // compiled code
+          export default {
+           title: 'Button',
+          }
+          export const Basic = () => React.createElement(Button);
+        `,
+          dedent`
+          // original code
           export default {
            title: 'Button',
           }
           /** The most basic button */
           export const Basic = () => <Button />
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
+        // compiled code
         export default {
           title: 'Button'
         };
-        /** The most basic button */
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -238,7 +322,16 @@ describe('enrichCsf', () => {
 
     it('JSDoc multi-line', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+        // compiled code
+        export default {
+         title: 'Button',
+        }
+        export const Basic = () => React.createElement(Button);
+      `,
+          dedent`
+          // original code
           export default {
            title: 'Button',
           }
@@ -248,17 +341,14 @@ describe('enrichCsf', () => {
            * In a block!
            */
           export const Basic = () => <Button />
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
+        // compiled code
         export default {
           title: 'Button'
         };
-        /**
-         * The most basic button
-         * 
-         * In a block!
-         */
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -278,27 +368,33 @@ describe('enrichCsf', () => {
 
     it('preserves indentation', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+          // compiled code
           export default {
            title: 'Button',
           }
-          /**
-           * - A bullet list
-           *   - A sub-bullet
-           * - A second bullet
-           */
-          export const Basic = () => <Button />
-        `)
-      ).toMatchInlineSnapshot(`
+          export const Basic = () => React.createElement(Button);
+        `,
+          dedent`
+        // original code
         export default {
-          title: 'Button'
-        };
+         title: 'Button',
+        }
         /**
          * - A bullet list
          *   - A sub-bullet
          * - A second bullet
          */
-        export const Basic = () => <Button />;
+        export const Basic = () => <Button />
+      `
+        )
+      ).toMatchInlineSnapshot(`
+        // compiled code
+        export default {
+          title: 'Button'
+        };
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -320,19 +416,29 @@ describe('enrichCsf', () => {
   describe('meta descriptions', () => {
     it('skips inline comments', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+        // compiled code
+        export default {
+          title: 'Button',
+        }
+        export const Basic = () => React.createElement(Button);
+        `,
+          dedent`
+        // original code
         // The most basic button
         export default {
-           title: 'Button',
-          }
-          export const Basic = () => <Button />
-        `)
+          title: 'Button',
+        }
+        export const Basic = () => <Button />
+        `
+        )
       ).toMatchInlineSnapshot(`
-        // The most basic button
+        // compiled code
         export default {
           title: 'Button'
         };
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -348,19 +454,29 @@ describe('enrichCsf', () => {
 
     it('skips blocks without jsdoc', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+        // compiled code
+        export default {
+         title: 'Button',
+        }
+        export const Basic = () => React.createElement();
+      `,
+          dedent`
+          // original code
           /* The most basic button */
           export default {
            title: 'Button',
           }
           export const Basic = () => <Button />
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
-        /* The most basic button */
+        // compiled code
         export default {
           title: 'Button'
         };
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement();
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -376,15 +492,25 @@ describe('enrichCsf', () => {
 
     it('JSDoc single-line', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+          // compiled code
+          export default {
+           title: 'Button'
+          }
+          export const Basic = () => React.createElement(Button)
+        `,
+          dedent`
+          // original code
           /** The most basic button */
           export default {
            title: 'Button'
           }
           export const Basic = () => <Button />
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
-        /** The most basic button */
+        // compiled code
         export default {
           title: 'Button',
           parameters: {
@@ -395,7 +521,7 @@ describe('enrichCsf', () => {
             }
           }
         };
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -411,7 +537,16 @@ describe('enrichCsf', () => {
 
     it('JSDoc multi-line', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+        // compiled code
+        export default {
+         title: 'Button',
+        }
+        export const Basic = () => React.createElement();
+      `,
+          dedent`
+          // original code
           /**
            * The most basic button
            * 
@@ -421,13 +556,10 @@ describe('enrichCsf', () => {
            title: 'Button',
           }
           export const Basic = () => <Button />
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
-        /**
-         * The most basic button
-         * 
-         * In a block!
-         */
+        // compiled code
         export default {
           title: 'Button',
           parameters: {
@@ -438,7 +570,7 @@ describe('enrichCsf', () => {
             }
           }
         };
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement();
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -454,7 +586,16 @@ describe('enrichCsf', () => {
 
     it('preserves indentation', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+        // compiled code
+        export default {
+         title: 'Button',
+        }
+        export const Basic = () => React.createElement(Button);
+      `,
+          dedent`
+          // original code
           /**
            * - A bullet list
            *   - A sub-bullet
@@ -464,13 +605,10 @@ describe('enrichCsf', () => {
            title: 'Button',
           }
           export const Basic = () => <Button />
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
-        /**
-         * - A bullet list
-         *   - A sub-bullet
-         * - A second bullet
-         */
+        // compiled code
         export default {
           title: 'Button',
           parameters: {
@@ -481,7 +619,7 @@ describe('enrichCsf', () => {
             }
           }
         };
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -497,7 +635,19 @@ describe('enrichCsf', () => {
 
     it('correctly interleaves parameters', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+        // compiled code
+        export default {
+          title: 'Button',
+          parameters: {
+            foo: 'bar',
+            docs: { inlineStories: true }
+          }
+        }
+        export const Basic = () => React.createElement(Button);
+      `,
+          dedent`
           /** The most basic button */
           export default {
             title: 'Button',
@@ -507,9 +657,10 @@ describe('enrichCsf', () => {
             }
           }
           export const Basic = () => <Button />
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
-        /** The most basic button */
+        // compiled code
         export default {
           title: 'Button',
           parameters: {
@@ -522,7 +673,7 @@ describe('enrichCsf', () => {
             }
           }
         };
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -538,7 +689,23 @@ describe('enrichCsf', () => {
 
     it('respects user component description', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+        // compiled code
+        export default {
+          title: 'Button',
+          parameters: {
+            docs: {
+              description: {
+                component: 'hahaha'
+              }
+            }
+          }
+        }
+        export const Basic = () => React.createElement(Button);
+      `,
+          dedent`
+          // original code
           /** The most basic button */
           export default {
             title: 'Button',
@@ -551,9 +718,10 @@ describe('enrichCsf', () => {
             }
           }
           export const Basic = () => <Button />
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
-        /** The most basic button */
+        // compiled code
         export default {
           title: 'Button',
           parameters: {
@@ -564,7 +732,7 @@ describe('enrichCsf', () => {
             }
           }
         };
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -580,7 +748,17 @@ describe('enrichCsf', () => {
 
     it('respects meta variables', () => {
       expect(
-        enrich(dedent`
+        enrich(
+          dedent`
+        // compiled code
+        const meta = {
+          title: 'Button'
+        }
+        export default meta;
+        export const Basic = () => React.createElement(Button);
+        `,
+          dedent`
+        // original code
         /** The most basic button */
         const meta = {
           title: 'Button'
@@ -588,9 +766,10 @@ describe('enrichCsf', () => {
         /** This should be ignored */
         export default meta;
         export const Basic = () => <Button />
-        `)
+        `
+        )
       ).toMatchInlineSnapshot(`
-        /** The most basic button */
+        // compiled code
         const meta = {
           title: 'Button',
           parameters: {
@@ -601,9 +780,8 @@ describe('enrichCsf', () => {
             }
           }
         };
-        /** This should be ignored */
         export default meta;
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -623,6 +801,14 @@ describe('enrichCsf', () => {
       expect(
         enrich(
           dedent`
+          // compiled code
+          export default {
+           title: 'Button',
+          }
+          export const Basic = () => React.createElement(Button);
+        `,
+          dedent`
+          // original code
           export default {
            title: 'Button',
           }
@@ -632,11 +818,11 @@ describe('enrichCsf', () => {
           { disableSource: true }
         )
       ).toMatchInlineSnapshot(`
+        // compiled code
         export default {
           title: 'Button'
         };
-        /** The most basic button */
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -654,6 +840,14 @@ describe('enrichCsf', () => {
       expect(
         enrich(
           dedent`
+          // compiled code
+          export default {
+           title: 'Button',
+          }
+          export const Basic = () => React.createElement(Button);
+        `,
+          dedent`
+          // original code
           export default {
            title: 'Button',
           }
@@ -663,11 +857,11 @@ describe('enrichCsf', () => {
           { disableDescription: true }
         )
       ).toMatchInlineSnapshot(`
+        // compiled code
         export default {
           title: 'Button'
         };
-        /** The most basic button */
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
         Basic.parameters = {
           ...Basic.parameters,
           docs: {
@@ -685,6 +879,14 @@ describe('enrichCsf', () => {
       expect(
         enrich(
           dedent`
+          // compiled code
+          export default {
+           title: 'Button',
+          }
+          export const Basic = () => React.createElement(Button);
+        `,
+          dedent`
+          // original code
           export default {
            title: 'Button',
           }
@@ -694,11 +896,11 @@ describe('enrichCsf', () => {
           { disableSource: true, disableDescription: true }
         )
       ).toMatchInlineSnapshot(`
+        // compiled code
         export default {
           title: 'Button'
         };
-        /** The most basic button */
-        export const Basic = () => <Button />;
+        export const Basic = () => React.createElement(Button);
       `);
     });
   });
diff --git a/code/lib/csf-tools/src/enrichCsf.ts b/code/lib/csf-tools/src/enrichCsf.ts
index dae96aa85c57..8744b8359eaa 100644
--- a/code/lib/csf-tools/src/enrichCsf.ts
+++ b/code/lib/csf-tools/src/enrichCsf.ts
@@ -9,10 +9,16 @@ export interface EnrichCsfOptions {
   disableDescription?: boolean;
 }
 
-export const enrichCsfStory = (csf: CsfFile, key: string, options?: EnrichCsfOptions) => {
-  const storyExport = csf.getStoryExport(key);
+export const enrichCsfStory = (
+  csf: CsfFile,
+  csfSource: CsfFile,
+  key: string,
+  options?: EnrichCsfOptions
+) => {
+  const storyExport = csfSource.getStoryExport(key);
   const source = !options?.disableSource && extractSource(storyExport);
-  const description = !options?.disableDescription && extractDescription(csf._storyStatements[key]);
+  const description =
+    !options?.disableDescription && extractDescription(csfSource._storyStatements[key]);
   const parameters = [];
   const originalParameters = t.memberExpression(t.identifier(key), t.identifier('parameters'));
   parameters.push(t.spreadElement(originalParameters));
@@ -111,8 +117,8 @@ const addComponentDescription = (
   addComponentDescription(subNode, rest, value);
 };
 
-export const enrichCsfMeta = (csf: CsfFile, options?: EnrichCsfOptions) => {
-  const description = !options?.disableDescription && extractDescription(csf._metaStatement);
+export const enrichCsfMeta = (csf: CsfFile, csfSource: CsfFile, options?: EnrichCsfOptions) => {
+  const description = !options?.disableDescription && extractDescription(csfSource._metaStatement);
   // docs: { description: { component: %%description%% } },
   if (description) {
     const metaNode = csf._metaNode;
@@ -126,10 +132,10 @@ export const enrichCsfMeta = (csf: CsfFile, options?: EnrichCsfOptions) => {
   }
 };
 
-export const enrichCsf = (csf: CsfFile, options?: EnrichCsfOptions) => {
-  enrichCsfMeta(csf, options);
+export const enrichCsf = (csf: CsfFile, csfSource: CsfFile, options?: EnrichCsfOptions) => {
+  enrichCsfMeta(csf, csfSource, options);
   Object.keys(csf._storyExports).forEach((key) => {
-    enrichCsfStory(csf, key, options);
+    enrichCsfStory(csf, csfSource, key, options);
   });
 };