diff --git a/__tests__/lib/infer/return.js b/__tests__/lib/infer/return.js index 979d640b0..6c7be8135 100644 --- 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 {}' + ); + 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 index c192307a5..0cb929194 100644 --- 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',