Add jest transform for hoisting module mocks (#540)

README.md

@@ -64,6 +64,9 @@ transforms are available:
   transform in the [react-hot-loader](https://github.com/gaearon/react-hot-loader)
   project. This enables advanced hot reloading use cases such as editing of
   bound methods.
+* **jest**: Hoist desired [jest](https://jestjs.io/) method calls above imports in
+  the same way as [babel-plugin-jest-hoist](https://github.com/facebook/jest/tree/master/packages/babel-plugin-jest-hoist).
+  Does not validate the arguments passed to `jest.mock`, but the same rules still apply.
 
 These proposed JS features are built-in and always transformed:
 * [Optional chaining](https://github.com/tc39/proposal-optional-chaining): `a?.b`

integrations/jest-plugin/README.md

@@ -21,5 +21,5 @@ Then change the default transform in jest.config.js file:
 ```
 
 Currently, the transforms are not configurable; it uses always runs the import
-transform and uses the file extension to decide whether to run the JSX, Flow,
-and/or TypeScript transforms.
+and jest transforms and uses the file extension to decide whether to run the
+JSX, Flow, and/or TypeScript transforms.

integrations/jest-plugin/src/index.ts

@@ -2,11 +2,11 @@ import {Transform, transform} from "sucrase";
 
 function getTransforms(filename: string): Array<Transform> | null {
   if (filename.endsWith(".js") || filename.endsWith(".jsx")) {
-    return ["flow", "jsx", "imports"];
+    return ["flow", "jsx", "imports", "jest"];
   } else if (filename.endsWith(".ts")) {
-    return ["typescript", "imports"];
+    return ["typescript", "imports", "jest"];
   } else if (filename.endsWith(".tsx")) {
-    return ["typescript", "jsx", "imports"];
+    return ["typescript", "jsx", "imports", "jest"];
   }
   return null;
 }

src/Options-gen-types.ts

@@ -10,6 +10,7 @@ export const Transform = t.union(
   t.lit("flow"),
   t.lit("imports"),
   t.lit("react-hot-loader"),
+  t.lit("jest"),
 );
 
 export const SourceMapOptions = t.iface([], {

src/Options.ts

@@ -4,7 +4,7 @@ import OptionsGenTypes from "./Options-gen-types";
 
 const {Options: OptionsChecker} = createCheckers(OptionsGenTypes);
 
-export type Transform = "jsx" | "typescript" | "flow" | "imports" | "react-hot-loader";
+export type Transform = "jsx" | "typescript" | "flow" | "imports" | "react-hot-loader" | "jest";
 
 export interface SourceMapOptions {
   /**

src/TokenProcessor.ts

@@ -172,7 +172,8 @@ export default class TokenProcessor {
   }
 
   removeToken(): void {
-    this.replaceTokenTrimmingLeftWhitespace("");
+    this.resultCode += this.previousWhitespaceAndComments().replace(/[^\r\n]/g, "");
+    this.tokenIndex++;
   }
 
   copyExpectedToken(tokenType: TokenType): void {

src/transformers/JestHoistTransformer.ts

@@ -0,0 +1,105 @@
+import type CJSImportProcessor from "../CJSImportProcessor";
+import {TokenType as tt} from "../parser/tokenizer/types";
+import type TokenProcessor from "../TokenProcessor";
+import type RootTransformer from "./RootTransformer";
+import Transformer from "./Transformer";
+
+const JEST_GLOBAL_NAME = "jest";
+const HOISTED_METHODS = ["mock", "unmock", "enableAutomock", "disableAutomock"];
+
+/**
+ * Implementation of babel-plugin-jest-hoist, which hoists up some jest method
+ * calls above the imports to allow them to override other imports.
+ */
+export default class JestHoistTransformer extends Transformer {
+  private readonly hoistedCalls: Array<string> = [];
+
+  constructor(
+    readonly rootTransformer: RootTransformer,
+    readonly tokens: TokenProcessor,
+    readonly importProcessor: CJSImportProcessor | null,
+  ) {
+    super();
+  }
+
+  process(): boolean {
+    if (
+      this.tokens.currentToken().scopeDepth === 0 &&
+      this.tokens.matches4(tt.name, tt.dot, tt.name, tt.parenL) &&
+      this.tokens.identifierName() === JEST_GLOBAL_NAME
+    ) {
+      // TODO: This only works if imports transform is active, which it will be for jest.
+      //       But if jest adds module support and we no longer need the import transform, this needs fixing.
+      if (this.importProcessor?.getGlobalNames()?.has(JEST_GLOBAL_NAME)) {
+        return false;
+      }
+      return this.extractHoistedCalls();
+    }
+
+    return false;
+  }
+
+  getHoistedCode(): string {
+    if (this.hoistedCalls.length > 0) {
+      // This will be placed before module interop code, but that's fine since
+      // imports aren't allowed in module mock factories.
+      return `\n${JEST_GLOBAL_NAME}${this.hoistedCalls.join("")};`;
+    }
+    return "";
+  }
+
+  /**
+   * Extracts any methods calls on the jest-object that should be hoisted.
+   *
+   * According to the jest docs, https://jestjs.io/docs/en/jest-object#jestmockmodulename-factory-options,
+   * mock, unmock, enableAutomock, disableAutomock, are the methods that should be hoisted.
+   *
+   * We do not apply the same checks of the arguments as babel-plugin-jest-hoist does.
+   */
+  private extractHoistedCalls(): boolean {
+    // We remove the `jest` expression, then add it back later if we find a non-hoisted call
+    this.tokens.removeToken();
+    let restoredJest = false;
+
+    // Iterate through all chained calls on the jest object
+    while (this.tokens.matches3(tt.dot, tt.name, tt.parenL)) {
+      const methodName = this.tokens.identifierNameAtIndex(this.tokens.currentIndex() + 1);
+      const shouldHoist = HOISTED_METHODS.includes(methodName);
+      if (shouldHoist) {
+        // We've matched e.g. `.mock(...)` or similar call
+        // Start by applying transforms to the entire call, including parameters
+        const snapshotBefore = this.tokens.snapshot();
+        this.tokens.copyToken();
+        this.tokens.copyToken();
+        this.tokens.copyToken();
+        this.rootTransformer.processBalancedCode();
+        this.tokens.copyExpectedToken(tt.parenR);
+        const snapshotAfter = this.tokens.snapshot();
+
+        // Then grab the transformed code and store it for hoisting
+        const processedCall = snapshotAfter.resultCode.slice(snapshotBefore.resultCode.length);
+        this.hoistedCalls.push(processedCall);
+
+        // Now go back and remove the entire method call
+        const endIndex = this.tokens.currentIndex();
+        this.tokens.restoreToSnapshot(snapshotBefore);
+        while (this.tokens.currentIndex() < endIndex) {
+          this.tokens.removeToken();
+        }
+      } else {
+        if (!restoredJest) {
+          restoredJest = true;
+          this.tokens.appendCode(JEST_GLOBAL_NAME);
+        }
+        // When not hoisting we just transform the method call as usual
+        this.tokens.copyToken();
+        this.tokens.copyToken();
+        this.tokens.copyToken();
+        this.rootTransformer.processBalancedCode();
+        this.tokens.copyExpectedToken(tt.parenR);
+      }
+    }
+
+    return true;
+  }
+}

src/transformers/RootTransformer.ts

@@ -8,6 +8,7 @@ import getClassInfo, {ClassInfo} from "../util/getClassInfo";
 import CJSImportTransformer from "./CJSImportTransformer";
 import ESMImportTransformer from "./ESMImportTransformer";
 import FlowTransformer from "./FlowTransformer";
+import JestHoistTransformer from "./JestHoistTransformer";
 import JSXTransformer from "./JSXTransformer";
 import NumericSeparatorTransformer from "./NumericSeparatorTransformer";
 import OptionalCatchBindingTransformer from "./OptionalCatchBindingTransformer";
@@ -100,6 +101,9 @@ export default class RootTransformer {
         new TypeScriptTransformer(this, tokenProcessor, transforms.includes("imports")),
       );
     }
+    if (transforms.includes("jest")) {
+      this.transformers.push(new JestHoistTransformer(this, tokenProcessor, importProcessor));
+    }
   }
 
   transform(): string {
@@ -113,6 +117,9 @@ export default class RootTransformer {
     }
     prefix += this.helperManager.emitHelpers();
     prefix += this.generatedVariables.map((v) => ` var ${v};`).join("");
+    for (const transformer of this.transformers) {
+      prefix += transformer.getHoistedCode();
+    }
     let suffix = "";
     for (const transformer of this.transformers) {
       suffix += transformer.getSuffixCode();

src/transformers/Transformer.ts

@@ -6,6 +6,10 @@ export default abstract class Transformer {
     return "";
   }
 
+  getHoistedCode(): string {
+    return "";
+  }
+
   getSuffixCode(): string {
     return "";
   }