feat: Add flow inference for generators

devongovett · tmcw · commit 7947e974d818 · 2019-04-18T09:47:51.000-07:00
diff --git a/__tests__/lib/infer/return.js b/__tests__/lib/infer/return.js
@@ -41,4 +41,26 @@ test('inferReturn', function() {
     name: 'string',
     type: 'NameExpression'
   });
+  const generatorFn = evaluate(
+    '/** */function *a(): Generator<Foo, Bar, Baz> {}'
+  );
+  expect(generatorFn.generator).toBe(true);
+  expect(generatorFn.yields).toEqual([
+    {
+      title: 'yields',
+      type: {
+        name: 'Foo',
+        type: 'NameExpression'
+      }
+    }
+  ]);
+  expect(generatorFn.returns).toEqual([
+    {
+      title: 'returns',
+      type: {
+        name: 'Bar',
+        type: 'NameExpression'
+      }
+    }
+  ]);
 });
diff --git a/src/infer/return.js b/src/infer/return.js
@@ -30,10 +30,27 @@ function inferReturn(comment) {
   }
 
   if (t.isFunction(fn) && fn.returnType && fn.returnType.typeAnnotation) {
-    const returnType = flowDoctrine(fn.returnType.typeAnnotation);
+    let returnType = flowDoctrine(fn.returnType.typeAnnotation);
     if (comment.returns && comment.returns.length > 0) {
       comment.returns[0].type = returnType;
     } else {
+      if (
+        fn.generator &&
+        returnType.type === 'TypeApplication' &&
+        returnType.expression.name === 'Generator' &&
+        returnType.applications.length === 3
+      ) {
+        comment.generator = true;
+        comment.yields = [
+          {
+            title: 'yields',
+            type: returnType.applications[0]
+          }
+        ];
+
+        returnType = returnType.applications[1];
+      }
+
       comment.returns = [
         {
           title: 'returns',
