Resolve flow $Keys to union when typeParameter is an ObjectTypeAnnotation (#290)

* Add support for ObjectTypeAnnotation as typeParameter of $Keys

* Add a test

* Refactor to use getPropertyName

* Update resolveObjectExpressionToNameArray to resolveObjectToNameArray

* Remove correctly failing $Keys test and add another $Keys test

* Re-add support for simple identifiers

Author: Jack Aldridge

src/utils/__tests__/getFlowType-test.js

@@ -536,6 +536,46 @@ describe('getFlowType', () => {
     });
   });
 
+  it('resolves $Keys with an ObjectTypeAnnotation typeParameter to union', () => {
+    const typePath = statement(`
+      var x: $Keys<{| apple: string, banana: string |}> = 2;
+    `)
+      .get('declarations', 0)
+      .get('id')
+      .get('typeAnnotation')
+      .get('typeAnnotation');
+
+    expect(getFlowType(typePath)).toEqual({
+      name: 'union',
+      elements: [
+        { name: 'literal', value: 'apple' },
+        { name: 'literal', value: 'banana' },
+      ],
+      raw: '$Keys<{| apple: string, banana: string |}>',
+    });
+  });
+
+  it('resolves $Keys with an ObjectTypeAnnotation typeParameter to union with an ObjectTypeSpreadProperty', () => {
+    const typePath = statement(`
+      var x: $Keys<{| apple: string, banana: string, ...OtherFruits |}> = 2;
+      type OtherFruits = { orange: string }
+    `)
+      .get('declarations', 0)
+      .get('id')
+      .get('typeAnnotation')
+      .get('typeAnnotation');
+
+    expect(getFlowType(typePath)).toEqual({
+      name: 'union',
+      elements: [
+        { name: 'literal', value: 'apple' },
+        { name: 'literal', value: 'banana' },
+        { name: 'literal', value: 'orange' },
+      ],
+      raw: '$Keys<{| apple: string, banana: string, ...OtherFruits |}>',
+    });
+  });
+
   it('handles multiple references to one type', () => {
     const typePath = statement(`
       let action: { a: Action, b: Action };

src/utils/getFlowType.js

@@ -17,7 +17,7 @@ import printValue from './printValue';
 import recast from 'recast';
 import getTypeAnnotation from '../utils/getTypeAnnotation';
 import resolveToValue from '../utils/resolveToValue';
-import { resolveObjectExpressionToNameArray } from '../utils/resolveObjectKeysToArray';
+import { resolveObjectToNameArray } from '../utils/resolveObjectKeysToArray';
 import type {
   FlowTypeDescriptor,
   FlowElementsType,
@@ -69,12 +69,16 @@ function handleKeysHelper(path: NodePath): ?FlowElementsType {
   let value = path.get('typeParameters', 'params', 0);
   if (types.TypeofTypeAnnotation.check(value.node)) {
     value = value.get('argument', 'id');
-  } else {
+  } else if (!types.ObjectTypeAnnotation.check(value.node)) {
     value = value.get('id');
   }
   const resolvedPath = resolveToValue(value);
-  if (resolvedPath && types.ObjectExpression.check(resolvedPath.node)) {
-    const keys = resolveObjectExpressionToNameArray(resolvedPath, true);
+  if (
+    resolvedPath &&
+    (types.ObjectExpression.check(resolvedPath.node) ||
+      types.ObjectTypeAnnotation.check(resolvedPath.node))
+  ) {
+    const keys = resolveObjectToNameArray(resolvedPath, true);
 
     if (keys) {
       return {

src/utils/resolveObjectKeysToArray.js

@@ -29,34 +29,57 @@ function isObjectKeysCall(node: ASTNode): boolean {
   );
 }
 
-export function resolveObjectExpressionToNameArray(
-  objectExpression: NodePath,
+function isWhitelistedObjectProperty(prop) {
+  return (
+    (types.Property.check(prop) &&
+      ((types.Identifier.check(prop.key) && !prop.computed) ||
+        types.Literal.check(prop.key))) ||
+    types.SpreadElement.check(prop)
+  );
+}
+
+function isWhiteListedObjectTypeProperty(prop) {
+  return (
+    types.ObjectTypeProperty.check(prop) ||
+    types.ObjectTypeSpreadProperty.check(prop)
+  );
+}
+
+// Resolves an ObjectExpression or an ObjectTypeAnnotation
+export function resolveObjectToNameArray(
+  object: NodePath,
   raw: boolean = false,
 ): ?Array<string> {
   if (
-    types.ObjectExpression.check(objectExpression.value) &&
-    objectExpression.value.properties.every(
-      prop =>
-        (types.Property.check(prop) &&
-          ((types.Identifier.check(prop.key) && !prop.computed) ||
-            types.Literal.check(prop.key))) ||
-        types.SpreadElement.check(prop),
-    )
+    (types.ObjectExpression.check(object.value) &&
+      object.value.properties.every(isWhitelistedObjectProperty)) ||
+    (types.ObjectTypeAnnotation.check(object.value) &&
+      object.value.properties.every(isWhiteListedObjectTypeProperty))
   ) {
     let values = [];
     let error = false;
-    objectExpression.get('properties').each(propPath => {
+    object.get('properties').each(propPath => {
       if (error) return;
       const prop = propPath.value;
 
-      if (types.Property.check(prop)) {
+      if (types.Property.check(prop) || types.ObjectTypeProperty.check(prop)) {
         // Key is either Identifier or Literal
         const name = prop.key.name || (raw ? prop.key.raw : prop.key.value);
 
         values.push(name);
-      } else if (types.SpreadElement.check(prop)) {
-        const spreadObject = resolveToValue(propPath.get('argument'));
-        const spreadValues = resolveObjectExpressionToNameArray(spreadObject);
+      } else if (
+        types.SpreadElement.check(prop) ||
+        types.ObjectTypeSpreadProperty.check(prop)
+      ) {
+        let spreadObject = resolveToValue(propPath.get('argument'));
+        if (types.GenericTypeAnnotation.check(spreadObject.value)) {
+          const typeAlias = resolveToValue(spreadObject.get('id'));
+          if (types.ObjectTypeAnnotation.check(typeAlias.get('right').value)) {
+            spreadObject = resolveToValue(typeAlias.get('right'));
+          }
+        }
+
+        const spreadValues = resolveObjectToNameArray(spreadObject);
         if (!spreadValues) {
           error = true;
           return;
@@ -87,7 +110,7 @@ export default function resolveObjectKeysToArray(path: NodePath): ?NodePath {
 
   if (isObjectKeysCall(node)) {
     const objectExpression = resolveToValue(path.get('arguments').get(0));
-    const values = resolveObjectExpressionToNameArray(objectExpression);
+    const values = resolveObjectToNameArray(objectExpression);
 
     if (values) {
       const nodes = values