fix(jest-snapshot): Fix a potential bug when using babel and improve …

…performance (#14036)
jestjs · Mar 31, 2023 · 257f250 · 257f250
1 parent b38d7d3
commit 257f250
Show file tree

Hide file tree

Showing 5 changed files with 121 additions and 123 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,7 @@
 
 - `[jest-environment-jsdom, jest-environment-node]` Fix assignment of `customExportConditions` via `testEnvironmentOptions` when custom env subclass defines a default value ([#13989](https://github.com/facebook/jest/pull/13989))
 - `[jest-matcher-utils]` Fix copying value of inherited getters ([#14007](https://github.com/facebook/jest/pull/14007))
+- `[jest-snapshot]` Fix a potential bug when not using prettier and improve performance ([#14036](https://github.com/facebook/jest/pull/14036))
 
 ### Chore & Maintenance
 

diff --git a/packages/jest-snapshot/package.json b/packages/jest-snapshot/package.json
@@ -21,12 +21,10 @@
     "@babel/generator": "^7.7.2",
     "@babel/plugin-syntax-jsx": "^7.7.2",
     "@babel/plugin-syntax-typescript": "^7.7.2",
-    "@babel/traverse": "^7.7.2",
     "@babel/types": "^7.3.3",
     "@jest/expect-utils": "workspace:^",
     "@jest/transform": "workspace:^",
     "@jest/types": "workspace:^",
-    "@types/babel__traverse": "^7.0.6",
     "@types/prettier": "^2.1.5",
     "babel-preset-current-node-syntax": "^1.0.0",
     "chalk": "^4.0.0",

diff --git a/packages/jest-snapshot/src/InlineSnapshots.ts b/packages/jest-snapshot/src/InlineSnapshots.ts
@@ -7,7 +7,13 @@
 
 import * as path from 'path';
 import type {ParseResult, PluginItem} from '@babel/core';
-import type {Expression, File, Program} from '@babel/types';
+import type {
+  Expression,
+  File,
+  Node,
+  Program,
+  TraversalAncestors,
+} from '@babel/types';
 import * as fs from 'graceful-fs';
 import type {
   CustomParser as PrettierCustomParser,
@@ -17,17 +23,18 @@ import semver = require('semver');
 import type {Frame} from 'jest-message-util';
 import {escapeBacktickString} from './utils';
 
-// prettier-ignore
-const babelTraverse = (
-  // @ts-expect-error requireOutside Babel transform
-  requireOutside('@babel/traverse') as typeof import('@babel/traverse')
-).default;
 // prettier-ignore
 const generate = (
   // @ts-expect-error requireOutside Babel transform
   requireOutside('@babel/generator') as typeof import('@babel/generator')
 ).default;
-const {file, isAwaitExpression, templateElement, templateLiteral} =
+const {
+  isAwaitExpression,
+  templateElement,
+  templateLiteral,
+  traverse,
+  traverseFast,
+} =
   // @ts-expect-error requireOutside Babel transform
   requireOutside('@babel/types') as typeof import('@babel/types');
 // @ts-expect-error requireOutside Babel transform
@@ -135,17 +142,22 @@ const saveSnapshotsForFile = (
   // substitute in the snapshots in reverse order, so slice calculations aren't thrown off.
   const sourceFileWithSnapshots = snapshots.reduceRight(
     (sourceSoFar, nextSnapshot) => {
+      const {node} = nextSnapshot;
       if (
-        !nextSnapshot.node ||
-        typeof nextSnapshot.node.start !== 'number' ||
-        typeof nextSnapshot.node.end !== 'number'
+        !node ||
+        typeof node.start !== 'number' ||
+        typeof node.end !== 'number'
       ) {
         throw new Error('Jest: no snapshot insert location found');
       }
+
+      // A hack to prevent unexpected line breaks in the generated code
+      node.loc!.end.line = node.loc!.start.line;
+
       return (
-        sourceSoFar.slice(0, nextSnapshot.node.start) +
-        generate(nextSnapshot.node, {retainLines: true}).code.trim() +
-        sourceSoFar.slice(nextSnapshot.node.end)
+        sourceSoFar.slice(0, node.start) +
+        generate(node, {retainLines: true}).code.trim() +
+        sourceSoFar.slice(node.end)
       );
     },
     sourceFile,
@@ -214,69 +226,56 @@ const indent = (snapshot: string, numIndents: number, indentation: string) => {
     .join('\n');
 };
 
-const resolveAst = (fileOrProgram: any): File => {
-  // Flow uses a 'Program' parent node, babel expects a 'File'.
-  let ast = fileOrProgram;
-  if (ast.type !== 'File') {
-    ast = file(ast, ast.comments, ast.tokens);
-    delete ast.program.comments;
-  }
-  return ast;
-};
-
 const traverseAst = (
   snapshots: Array<InlineSnapshot>,
-  fileOrProgram: File | Program,
+  ast: File | Program,
   snapshotMatcherNames: Array<string>,
 ) => {
-  const ast = resolveAst(fileOrProgram);
   const groupedSnapshots = groupSnapshotsByFrame(snapshots);
   const remainingSnapshots = new Set(snapshots.map(({snapshot}) => snapshot));
 
-  babelTraverse(ast, {
-    CallExpression({node}) {
-      const {arguments: args, callee} = node;
-      if (
-        callee.type !== 'MemberExpression' ||
-        callee.property.type !== 'Identifier' ||
-        callee.property.loc == null
-      ) {
-        return;
-      }
-      const {line, column} = callee.property.loc.start;
-      const snapshotsForFrame = groupedSnapshots[`${line}:${column}`];
-      if (!snapshotsForFrame) {
-        return;
-      }
-      if (snapshotsForFrame.length > 1) {
-        throw new Error(
-          'Jest: Multiple inline snapshots for the same call are not supported.',
-        );
-      }
-
-      snapshotMatcherNames.push(callee.property.name);
+  traverseFast(ast, (node: Node) => {
+    if (node.type !== 'CallExpression') return;
 
-      const snapshotIndex = args.findIndex(
-        ({type}) => type === 'TemplateLiteral',
+    const {arguments: args, callee} = node;
+    if (
+      callee.type !== 'MemberExpression' ||
+      callee.property.type !== 'Identifier' ||
+      callee.property.loc == null
+    ) {
+      return;
+    }
+    const {line, column} = callee.property.loc.start;
+    const snapshotsForFrame = groupedSnapshots[`${line}:${column}`];
+    if (!snapshotsForFrame) {
+      return;
+    }
+    if (snapshotsForFrame.length > 1) {
+      throw new Error(
+        'Jest: Multiple inline snapshots for the same call are not supported.',
       );
-      const values = snapshotsForFrame.map(inlineSnapshot => {
-        inlineSnapshot.node = node;
-        const {snapshot} = inlineSnapshot;
-        remainingSnapshots.delete(snapshot);
-
-        return templateLiteral(
-          [templateElement({raw: escapeBacktickString(snapshot)})],
-          [],
-        );
-      });
-      const replacementNode = values[0];
-
-      if (snapshotIndex > -1) {
-        args[snapshotIndex] = replacementNode;
-      } else {
-        args.push(replacementNode);
-      }
-    },
+    }
+    const inlineSnapshot = snapshotsForFrame[0];
+    inlineSnapshot.node = node;
+
+    snapshotMatcherNames.push(callee.property.name);
+
+    const snapshotIndex = args.findIndex(
+      ({type}) => type === 'TemplateLiteral',
+    );
+
+    const {snapshot} = inlineSnapshot;
+    remainingSnapshots.delete(snapshot);
+    const replacementNode = templateLiteral(
+      [templateElement({raw: escapeBacktickString(snapshot)})],
+      [],
+    );
+
+    if (snapshotIndex > -1) {
+      args[snapshotIndex] = replacementNode;
+    } else {
+      args.push(replacementNode);
+    }
   });
 
   if (remainingSnapshots.size) {
@@ -340,59 +339,61 @@ const createFormattingParser =
     // Workaround for https://github.com/prettier/prettier/issues/3150
     options.parser = inferredParser;
 
-    const ast = resolveAst(parsers[inferredParser](text, options));
-    babelTraverse(ast, {
-      CallExpression({node: {arguments: args, callee}, parent}) {
-        if (
-          callee.type !== 'MemberExpression' ||
-          callee.property.type !== 'Identifier' ||
-          !snapshotMatcherNames.includes(callee.property.name) ||
-          !callee.loc ||
-          callee.computed
-        ) {
-          return;
-        }
+    const ast = parsers[inferredParser](text, options);
+    traverse(ast, (node: Node, ancestors: TraversalAncestors) => {
+      if (node.type !== 'CallExpression') return;
 
-        let snapshotIndex: number | undefined;
-        let snapshot: string | undefined;
-        for (let i = 0; i < args.length; i++) {
-          const node = args[i];
-          if (node.type === 'TemplateLiteral') {
-            snapshotIndex = i;
-            snapshot = node.quasis[0].value.raw;
-          }
-        }
-        if (snapshot === undefined || snapshotIndex === undefined) {
-          return;
+      const {arguments: args, callee} = node;
+      if (
+        callee.type !== 'MemberExpression' ||
+        callee.property.type !== 'Identifier' ||
+        !snapshotMatcherNames.includes(callee.property.name) ||
+        !callee.loc ||
+        callee.computed
+      ) {
+        return;
+      }
+
+      let snapshotIndex: number | undefined;
+      let snapshot: string | undefined;
+      for (let i = 0; i < args.length; i++) {
+        const node = args[i];
+        if (node.type === 'TemplateLiteral') {
+          snapshotIndex = i;
+          snapshot = node.quasis[0].value.raw;
         }
+      }
+      if (snapshot === undefined) {
+        return;
+      }
 
-        const startColumn =
-          isAwaitExpression(parent) && parent.loc
-            ? parent.loc.start.column
-            : callee.loc.start.column;
-
-        const useSpaces = !options.useTabs;
-        snapshot = indent(
-          snapshot,
-          Math.ceil(
-            useSpaces
-              ? startColumn / (options.tabWidth ?? 1)
-              : // Each tab is 2 characters.
-                startColumn / 2,
-          ),
-          useSpaces ? ' '.repeat(options.tabWidth ?? 1) : '\t',
-        );
-
-        const replacementNode = templateLiteral(
-          [
-            templateElement({
-              raw: snapshot,
-            }),
-          ],
-          [],
-        );
-        args[snapshotIndex] = replacementNode;
-      },
+      const parent = ancestors[ancestors.length - 1].node;
+      const startColumn =
+        isAwaitExpression(parent) && parent.loc
+          ? parent.loc.start.column
+          : callee.loc.start.column;
+
+      const useSpaces = !options.useTabs;
+      snapshot = indent(
+        snapshot,
+        Math.ceil(
+          useSpaces
+            ? startColumn / (options.tabWidth ?? 1)
+            : // Each tab is 2 characters.
+              startColumn / 2,
+        ),
+        useSpaces ? ' '.repeat(options.tabWidth ?? 1) : '\t',
+      );
+
+      const replacementNode = templateLiteral(
+        [
+          templateElement({
+            raw: snapshot,
+          }),
+        ],
+        [],
+      );
+      args[snapshotIndex!] = replacementNode;
     });
 
     return ast;

diff --git a/packages/jest-snapshot/src/utils.ts b/packages/jest-snapshot/src/utils.ts
@@ -194,7 +194,7 @@ const printBacktickString = (str: string): string =>
 
 export const ensureDirectoryExists = (filePath: string): void => {
   try {
-    fs.mkdirSync(path.join(path.dirname(filePath)), {recursive: true});
+    fs.mkdirSync(path.dirname(filePath), {recursive: true});
   } catch {}
 };
 

diff --git a/yarn.lock b/yarn.lock
@@ -1832,7 +1832,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@babel/traverse@npm:^7.12.9, @babel/traverse@npm:^7.18.8, @babel/traverse@npm:^7.20.0, @babel/traverse@npm:^7.20.10, @babel/traverse@npm:^7.20.12, @babel/traverse@npm:^7.20.13, @babel/traverse@npm:^7.20.5, @babel/traverse@npm:^7.20.7, @babel/traverse@npm:^7.7.2":
+"@babel/traverse@npm:^7.12.9, @babel/traverse@npm:^7.18.8, @babel/traverse@npm:^7.20.0, @babel/traverse@npm:^7.20.10, @babel/traverse@npm:^7.20.12, @babel/traverse@npm:^7.20.13, @babel/traverse@npm:^7.20.5, @babel/traverse@npm:^7.20.7":
   version: 7.20.13
   resolution: "@babel/traverse@npm:7.20.13"
   dependencies:
@@ -13096,15 +13096,13 @@ __metadata:
     "@babel/plugin-syntax-typescript": ^7.7.2
     "@babel/preset-flow": ^7.7.2
     "@babel/preset-react": ^7.12.1
-    "@babel/traverse": ^7.7.2
     "@babel/types": ^7.3.3
     "@jest/expect-utils": "workspace:^"
     "@jest/test-utils": "workspace:^"
     "@jest/transform": "workspace:^"
     "@jest/types": "workspace:^"
     "@tsd/typescript": ^4.9.0
     "@types/babel__core": ^7.1.14
-    "@types/babel__traverse": ^7.0.6
     "@types/graceful-fs": ^4.1.3
     "@types/natural-compare": ^1.4.0
     "@types/prettier": ^2.1.5