Skip to content

Commit 7947e97

Browse files
devongovetttmcw
authored andcommitted
feat: Add flow inference for generators
1 parent 8e3cd47 commit 7947e97

File tree

2 files changed

+40
-1
lines changed

2 files changed

+40
-1
lines changed

__tests__/lib/infer/return.js

+22
Original file line numberDiff line numberDiff line change
@@ -41,4 +41,26 @@ test('inferReturn', function() {
4141
name: 'string',
4242
type: 'NameExpression'
4343
});
44+
const generatorFn = evaluate(
45+
'/** */function *a(): Generator<Foo, Bar, Baz> {}'
46+
);
47+
expect(generatorFn.generator).toBe(true);
48+
expect(generatorFn.yields).toEqual([
49+
{
50+
title: 'yields',
51+
type: {
52+
name: 'Foo',
53+
type: 'NameExpression'
54+
}
55+
}
56+
]);
57+
expect(generatorFn.returns).toEqual([
58+
{
59+
title: 'returns',
60+
type: {
61+
name: 'Bar',
62+
type: 'NameExpression'
63+
}
64+
}
65+
]);
4466
});

src/infer/return.js

+18-1
Original file line numberDiff line numberDiff line change
@@ -30,10 +30,27 @@ function inferReturn(comment) {
3030
}
3131

3232
if (t.isFunction(fn) && fn.returnType && fn.returnType.typeAnnotation) {
33-
const returnType = flowDoctrine(fn.returnType.typeAnnotation);
33+
let returnType = flowDoctrine(fn.returnType.typeAnnotation);
3434
if (comment.returns && comment.returns.length > 0) {
3535
comment.returns[0].type = returnType;
3636
} else {
37+
if (
38+
fn.generator &&
39+
returnType.type === 'TypeApplication' &&
40+
returnType.expression.name === 'Generator' &&
41+
returnType.applications.length === 3
42+
) {
43+
comment.generator = true;
44+
comment.yields = [
45+
{
46+
title: 'yields',
47+
type: returnType.applications[0]
48+
}
49+
];
50+
51+
returnType = returnType.applications[1];
52+
}
53+
3754
comment.returns = [
3855
{
3956
title: 'returns',

0 commit comments

Comments
 (0)