From b4f20569529d6dd361713490dc58904949e564d3 Mon Sep 17 00:00:00 2001
From: Michael Shilman <shilman@users.noreply.github.com>
Date: Sun, 3 Apr 2022 07:48:16 +0800
Subject: [PATCH] Merge pull request #17858 from
 storybookjs/tech/automigrate-preserve-quotes

CLI: Preserve quote style in automigrate
---
 lib/csf-tools/src/ConfigFile.test.ts | 39 ++++++++++++++---
 lib/csf-tools/src/ConfigFile.ts      | 64 ++++++++++++++++++++--------
 lib/csf-tools/src/babelParse.ts      |  1 +
 3 files changed, 81 insertions(+), 23 deletions(-)

diff --git a/lib/csf-tools/src/ConfigFile.test.ts b/lib/csf-tools/src/ConfigFile.test.ts
index cf12a3dbb5fb..76b6ccb38d55 100644
--- a/lib/csf-tools/src/ConfigFile.test.ts
+++ b/lib/csf-tools/src/ConfigFile.test.ts
@@ -174,7 +174,7 @@ describe('ConfigFile', () => {
         ).toMatchInlineSnapshot(`
           export const core = {
             foo: 'bar',
-            builder: "webpack5"
+            builder: 'webpack5'
           };
         `);
       });
@@ -189,7 +189,7 @@ describe('ConfigFile', () => {
           )
         ).toMatchInlineSnapshot(`
           export const core = {
-            builder: "webpack5"
+            builder: 'webpack5'
           };
         `);
       });
@@ -205,7 +205,7 @@ describe('ConfigFile', () => {
         ).toMatchInlineSnapshot(`
           export const core = {
             builder: {
-              "name": "webpack5"
+              name: 'webpack5'
             }
           };
         `);
@@ -222,7 +222,7 @@ describe('ConfigFile', () => {
           )
         ).toMatchInlineSnapshot(`
           const coreVar = {
-            builder: "webpack5"
+            builder: 'webpack5'
           };
           export const core = coreVar;
         `);
@@ -260,7 +260,7 @@ describe('ConfigFile', () => {
           module.exports = {
             core: {
               foo: 'bar',
-              builder: "webpack5"
+              builder: 'webpack5'
             }
           };
         `);
@@ -277,11 +277,38 @@ describe('ConfigFile', () => {
         ).toMatchInlineSnapshot(`
           module.exports = {
             core: {
-              builder: "webpack5"
+              builder: 'webpack5'
             }
           };
         `);
       });
     });
+    describe('quotes', () => {
+      it('no quotes', () => {
+        expect(setField(['foo', 'bar'], 'baz', '')).toMatchInlineSnapshot(`
+          export const foo = {
+            bar: "baz"
+          };
+        `);
+      });
+      it('more single quotes', () => {
+        expect(setField(['foo', 'bar'], 'baz', `export const stories = ['a', 'b', "c"]`))
+          .toMatchInlineSnapshot(`
+          export const stories = ['a', 'b', "c"];
+          export const foo = {
+            bar: 'baz'
+          };
+        `);
+      });
+      it('more double quotes', () => {
+        expect(setField(['foo', 'bar'], 'baz', `export const stories = ['a', "b", "c"]`))
+          .toMatchInlineSnapshot(`
+          export const stories = ['a', "b", "c"];
+          export const foo = {
+            bar: "baz"
+          };
+        `);
+      });
+    });
   });
 });
diff --git a/lib/csf-tools/src/ConfigFile.ts b/lib/csf-tools/src/ConfigFile.ts
index c72e1d4e0117..9e67096e307e 100644
--- a/lib/csf-tools/src/ConfigFile.ts
+++ b/lib/csf-tools/src/ConfigFile.ts
@@ -81,14 +81,19 @@ const _updateExportNode = (path: string[], expr: t.Expression, existing: t.Objec
 export class ConfigFile {
   _ast: t.File;
 
+  _code: string;
+
   _exports: Record<string, t.Expression> = {};
 
   _exportsObject: t.ObjectExpression;
 
+  _quotes: 'single' | 'double' | undefined;
+
   fileName?: string;
 
-  constructor(ast: t.File, fileName?: string) {
+  constructor(ast: t.File, code: string, fileName?: string) {
     this._ast = ast;
+    this._code = code;
     this.fileName = fileName;
   }
 
@@ -190,24 +195,49 @@ export class ConfigFile {
     }
   }
 
+  _inferQuotes() {
+    if (!this._quotes) {
+      // first 500 tokens for efficiency
+      const occurrences = (this._ast.tokens || []).slice(0, 500).reduce(
+        (acc, token) => {
+          if (token.type.label === 'string') {
+            acc[this._code[token.start]] += 1;
+          }
+          return acc;
+        },
+        { "'": 0, '"': 0 }
+      );
+      this._quotes = occurrences["'"] > occurrences['"'] ? 'single' : 'double';
+    }
+    return this._quotes;
+  }
+
   setFieldValue(path: string[], value: any) {
-    const stringified = JSON.stringify(value);
-    const program = babelParse(`const __x = ${stringified}`);
+    const quotes = this._inferQuotes();
     let valueNode;
-    traverse(program, {
-      VariableDeclaration: {
-        enter({ node }) {
-          if (
-            node.declarations.length === 1 &&
-            t.isVariableDeclarator(node.declarations[0]) &&
-            t.isIdentifier(node.declarations[0].id) &&
-            node.declarations[0].id.name === '__x'
-          ) {
-            valueNode = node.declarations[0].init;
-          }
+    // we do this rather than t.valueToNode because apparently
+    // babel only preserves quotes if they are parsed from the original code.
+    if (quotes === 'single') {
+      const { code } = generate(t.valueToNode(value), { jsescOption: { quotes } });
+      const program = babelParse(`const __x = ${code}`);
+      traverse(program, {
+        VariableDeclaration: {
+          enter({ node }) {
+            if (
+              node.declarations.length === 1 &&
+              t.isVariableDeclarator(node.declarations[0]) &&
+              t.isIdentifier(node.declarations[0].id) &&
+              node.declarations[0].id.name === '__x'
+            ) {
+              valueNode = node.declarations[0].init;
+            }
+          },
         },
-      },
-    });
+      });
+    } else {
+      // double quotes is the default so we can skip all that
+      valueNode = t.valueToNode(value);
+    }
     if (!valueNode) {
       throw new Error(`Unexpected value ${JSON.stringify(value)}`);
     }
@@ -217,7 +247,7 @@ export class ConfigFile {
 
 export const loadConfig = (code: string, fileName?: string) => {
   const ast = babelParse(code);
-  return new ConfigFile(ast, fileName);
+  return new ConfigFile(ast, code, fileName);
 };
 
 export const formatConfig = (config: ConfigFile) => {
diff --git a/lib/csf-tools/src/babelParse.ts b/lib/csf-tools/src/babelParse.ts
index 406d58a96990..70c85fe17489 100644
--- a/lib/csf-tools/src/babelParse.ts
+++ b/lib/csf-tools/src/babelParse.ts
@@ -10,4 +10,5 @@ export const babelParse = (code: string) =>
       ['decorators', { decoratorsBeforeExport: true }],
       'classProperties',
     ],
+    tokens: true,
   });