Support some TSTypes in the inferrer

* add ts inferer tests
* fix: support some TSTypes in inferrer
* Update packages/babel-plugin-transform-for-of/src/index.ts
* fix typing errors
* fix(inferer): allow {} prototype name in generics

Co-authored-by: Nicolò Ribaudo <nicolo.ribaudo@gmail.com>

packages/babel-plugin-transform-for-of/src/index.ts

@@ -208,8 +208,10 @@ export default declare((api, options: Options) => {
         const right = path.get("right");
         if (
           right.isArrayExpression() ||
-          right.isGenericType("Array") ||
-          t.isArrayTypeAnnotation(right.getTypeAnnotation())
+          (process.env.BABEL_8_BREAKING
+            ? right.isGenericType("Array")
+            : right.isGenericType("Array") ||
+              t.isArrayTypeAnnotation(right.getTypeAnnotation()))
         ) {
           path.replaceWith(_ForOfStatementArray(path));
           return;

packages/babel-plugin-transform-for-of/test/fixtures/opt/options.json

@@ -1,3 +1,4 @@
 {
-  "plugins": ["transform-for-of", "transform-flow-strip-types"]
+  "plugins": ["transform-for-of"],
+  "presets": ["typescript"]
 }

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-flow-deopt/input.js

@@ -0,0 +1,4 @@
+// This won't be optimize because when `for-of` is handled, the b's type annotation has been removed by the Flow plugin
+function a(b: Array<any>) {
+  for (const y of b) {}
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-flow-deopt/options.json

@@ -0,0 +1,3 @@
+{
+  "plugins": ["transform-for-of", "transform-flow-strip-types"]
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-flow-deopt/output.js

@@ -0,0 +1,15 @@
+// This won't be optimize because when `for-of` is handled, the b's type annotation has been removed by the Flow plugin
+function a(b) {
+  var _iterator = babelHelpers.createForOfIteratorHelper(b),
+      _step;
+
+  try {
+    for (_iterator.s(); !(_step = _iterator.n()).done;) {
+      const y = _step.value;
+    }
+  } catch (err) {
+    _iterator.e(err);
+  } finally {
+    _iterator.f();
+  }
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-flow/input.js

@@ -0,0 +1 @@
+for (const y of (b: Array<any>)) {}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-flow/options.json

@@ -0,0 +1,3 @@
+{
+  "plugins": ["transform-for-of", "transform-flow-strip-types"]
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-flow/output.js

@@ -0,0 +1,3 @@
+for (var _i = 0, _arr = b; _i < _arr.length; _i++) {
+  const y = _arr[_i];
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-ts-deopt/input.ts

@@ -0,0 +1,4 @@
+// This won't be optimize because when `for-of` is handled, the b's type annotation has been removed by the TS plugin
+function a(b: Array<any>) {
+  for (const y of b) {}
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-ts-deopt/options.json

@@ -0,0 +1,4 @@
+{
+  "plugins": ["transform-for-of"],
+  "presets": ["typescript"]
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/typeannotation/output.js → packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-ts-deopt/output.js

@@ -1,7 +1,4 @@
-for (var _i = 0, _arr = b; _i < _arr.length; _i++) {
-  const y = _arr[_i];
-}
-
+// This won't be optimize because when `for-of` is handled, the b's type annotation has been removed by the TS plugin
 function a(b) {
   var _iterator = babelHelpers.createForOfIteratorHelper(b),
       _step;

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-ts/input.ts

@@ -0,0 +1 @@
+for (const y of (b as Array<any>)) {}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-ts/options.json

@@ -0,0 +1,4 @@
+{
+  "plugins": ["transform-for-of"],
+  "presets": ["typescript"]
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-ts/output.js

@@ -0,0 +1,3 @@
+for (var _i = 0, _arr = b; _i < _arr.length; _i++) {
+  const y = _arr[_i];
+}

packages/babel-plugin-transform-spread/test/fixtures/spread/type-reference-flow/input.ts

@@ -0,0 +1,14 @@
+function foo() {
+  const x = 1 ? a() : b();
+  const y = a() || b();
+
+  return [...x, ...y];
+}
+
+function a(): Array<number> {
+  return [];
+}
+
+function b(): Array<number> {
+  return [];
+}

packages/babel-plugin-transform-spread/test/fixtures/spread/type-reference-flow/options.json

@@ -0,0 +1,4 @@
+{
+  "plugins": ["transform-spread"],
+  "presets": ["flow"]
+}

packages/babel-plugin-transform-spread/test/fixtures/spread/type-reference-flow/output.js

@@ -0,0 +1,13 @@
+function foo() {
+  const x = 1 ? a() : b();
+  const y = a() || b();
+  return [].concat(x, y);
+}
+
+function a() {
+  return [];
+}
+
+function b() {
+  return [];
+}

packages/babel-plugin-transform-spread/test/fixtures/spread/type-reference-ts/input.ts

@@ -0,0 +1,14 @@
+function foo() {
+  const x = 1 ? a() : b();
+  const y = a() || b();
+
+  return [...x, ...y];
+}
+
+function a(): Array<number> {
+  return [];
+}
+
+function b(): Array<number> {
+  return [];
+}

packages/babel-plugin-transform-spread/test/fixtures/spread/type-reference-ts/output.js

@@ -0,0 +1,13 @@
+function foo() {
+  const x = 1 ? a() : b();
+  const y = a() || b();
+  return [].concat(x, y);
+}
+
+function a() {
+  return [];
+}
+
+function b() {
+  return [];
+}

packages/babel-traverse/src/path/inference/index.ts

@@ -3,6 +3,7 @@ import * as inferers from "./inferers";
 import {
   anyTypeAnnotation,
   isAnyTypeAnnotation,
+  isArrayTypeAnnotation,
   isBooleanTypeAnnotation,
   isEmptyTypeAnnotation,
   isFlowBaseAnnotation,
@@ -11,6 +12,10 @@ import {
   isMixedTypeAnnotation,
   isNumberTypeAnnotation,
   isStringTypeAnnotation,
+  isTSArrayType,
+  isTSTypeAnnotation,
+  isTSTypeReference,
+  isTupleTypeAnnotation,
   isTypeAnnotation,
   isUnionTypeAnnotation,
   isVoidTypeAnnotation,
@@ -23,13 +28,15 @@ import type * as t from "@babel/types";
  * Infer the type of the current `NodePath`.
  */
 
-export function getTypeAnnotation(this: NodePath): t.FlowType {
+export function getTypeAnnotation(this: NodePath): t.FlowType | t.TSType {
   let type = this.getData("typeAnnotation");
   if (type != null) {
     return type;
   }
   type = this._getTypeAnnotation() || anyTypeAnnotation();
-  if (isTypeAnnotation(type)) type = type.typeAnnotation;
+  if (isTypeAnnotation(type) || isTSTypeAnnotation(type)) {
+    type = type.typeAnnotation;
+  }
   this.setData("typeAnnotation", type);
   return type;
 }
@@ -108,7 +115,7 @@ export function isBaseType(
 
 function _isBaseType(
   baseName: string,
-  type?: t.FlowType,
+  type?: t.FlowType | t.TSType,
   soft?: boolean,
 ): boolean {
   if (baseName === "string") {
@@ -165,8 +172,24 @@ export function baseTypeStrictlyMatches(
 
 export function isGenericType(this: NodePath, genericName: string): boolean {
   const type = this.getTypeAnnotation();
+  if (genericName === "Array") {
+    // T[]
+    if (
+      isTSArrayType(type) ||
+      isArrayTypeAnnotation(type) ||
+      isTupleTypeAnnotation(type)
+    ) {
+      return true;
+    }
+  }
   return (
-    isGenericTypeAnnotation(type) &&
-    isIdentifier(type.id, { name: genericName })
+    (isGenericTypeAnnotation(type) &&
+      isIdentifier(type.id, {
+        name: genericName,
+      })) ||
+    (isTSTypeReference(type) &&
+      isIdentifier(type.typeName, {
+        name: genericName,
+      }))
   );
 }

packages/babel-traverse/src/path/inference/inferer-reference.ts

@@ -1,17 +1,15 @@
 import type NodePath from "../index";
 import {
   BOOLEAN_NUMBER_BINARY_OPERATORS,
-  createFlowUnionType,
-  createTSUnionType,
   createTypeAnnotationBasedOnTypeof,
-  createUnionTypeAnnotation,
-  isTSTypeAnnotation,
   numberTypeAnnotation,
   voidTypeAnnotation,
 } from "@babel/types";
 import type * as t from "@babel/types";
 import type Binding from "../../scope/binding";
 
+import { createUnionType } from "./util";
+
 export default function (this: NodePath<t.Identifier>, node: t.Identifier) {
   if (!this.isReferenced()) return;
 
@@ -110,18 +108,7 @@ function getTypeAnnotationBindingConstantViolations(
     return;
   }
 
-  if (isTSTypeAnnotation(types[0]) && createTSUnionType) {
-    return createTSUnionType(
-      // @ts-ignore fixme: createTSUnionType should handle voidTypeAnnotation
-      types as t.TSTypeAnnotation[],
-    );
-  }
-
-  if (createFlowUnionType) {
-    return createFlowUnionType(types as t.FlowType[]);
-  }
-
-  return createUnionTypeAnnotation(types as t.FlowType[]);
+  return createUnionType(types);
 }
 
 function getConstantViolationsBefore(
@@ -249,25 +236,8 @@ function getConditionalAnnotation<T extends t.Node>(
   }
 
   if (types.length) {
-    if (isTSTypeAnnotation(types[0]) && createTSUnionType) {
-      return {
-        typeAnnotation: createTSUnionType(
-          // @ts-ignore fixme: createTSUnionType should handle voidTypeAnnotation
-          types as t.TSTypeAnnotation[],
-        ),
-        ifStatement,
-      };
-    }
-
-    if (createFlowUnionType) {
-      return {
-        typeAnnotation: createFlowUnionType(types),
-        ifStatement,
-      };
-    }
-
     return {
-      typeAnnotation: createUnionTypeAnnotation(types),
+      typeAnnotation: createUnionType(types),
       ifStatement,
     };
   }