add tests and initial fix

Author: gabritto

src/compiler/checker.ts

@@ -26672,11 +26672,22 @@ namespace ts {
             const func = getContainingFunction(node);
             if (func) {
                 const functionFlags = getFunctionFlags(func);
-                const contextualReturnType = getContextualReturnType(func);
+                let contextualReturnType = getContextualReturnType(func);
                 if (contextualReturnType) {
+                    const isAsyncGenerator = (functionFlags & FunctionFlags.Async) !== 0;
+                    if (!node.asteriskToken && contextualReturnType.flags & TypeFlags.Union) {
+                        contextualReturnType = filterType(contextualReturnType, type => {
+                            // >> What should we use to filter?
+                            // >> What will be the perf impacts? Is there caching? Is it enough to make this ok?
+                            // return !!getIterationTypeOfGeneratorFunctionReturnType(IterationTypeKind.Yield, type, isAsyncGenerator);
+                            // >> ^^^ doesn't work, because we get { yield: any, return: any } for bogus types
+                            const generator = createGeneratorReturnType(anyType, anyType, anyType, isAsyncGenerator);
+                            return checkTypeAssignableTo(generator, type, /*errorNode*/ undefined);
+                        });
+                    } // >> TODO: should we do something when `node.asteriskToken`?
                     return node.asteriskToken
                         ? contextualReturnType
-                        : getIterationTypeOfGeneratorFunctionReturnType(IterationTypeKind.Yield, contextualReturnType, (functionFlags & FunctionFlags.Async) !== 0);
+                        : getIterationTypeOfGeneratorFunctionReturnType(IterationTypeKind.Yield, contextualReturnType, isAsyncGenerator);
                 }
             }
 

tests/baselines/reference/contextualTypeOnYield.js

@@ -0,0 +1,12 @@
+//// [contextualTypeOnYield.ts]
+type FuncOrGeneratorFunc = () => (number | Generator<(arg: number) => void, any, void>)
+
+const f: FuncOrGeneratorFunc = function*() {
+  yield (num) => console.log(num); // `num` should be inferred to have type `number`.
+}
+
+//// [contextualTypeOnYield.js]
+"use strict";
+const f = function* () {
+    yield (num) => console.log(num); // `num` should be inferred to have type `number`.
+};

tests/baselines/reference/contextualTypeOnYield.symbols

@@ -0,0 +1,17 @@
+=== tests/cases/compiler/contextualTypeOnYield.ts ===
+type FuncOrGeneratorFunc = () => (number | Generator<(arg: number) => void, any, void>)
+>FuncOrGeneratorFunc : Symbol(FuncOrGeneratorFunc, Decl(contextualTypeOnYield.ts, 0, 0))
+>Generator : Symbol(Generator, Decl(lib.es2015.generator.d.ts, --, --))
+>arg : Symbol(arg, Decl(contextualTypeOnYield.ts, 0, 54))
+
+const f: FuncOrGeneratorFunc = function*() {
+>f : Symbol(f, Decl(contextualTypeOnYield.ts, 2, 5))
+>FuncOrGeneratorFunc : Symbol(FuncOrGeneratorFunc, Decl(contextualTypeOnYield.ts, 0, 0))
+
+  yield (num) => console.log(num); // `num` should be inferred to have type `number`.
+>num : Symbol(num, Decl(contextualTypeOnYield.ts, 3, 9))
+>console.log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
+>console : Symbol(console, Decl(lib.dom.d.ts, --, --))
+>log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
+>num : Symbol(num, Decl(contextualTypeOnYield.ts, 3, 9))
+}

tests/baselines/reference/contextualTypeOnYield.types

@@ -0,0 +1,19 @@
+=== tests/cases/compiler/contextualTypeOnYield.ts ===
+type FuncOrGeneratorFunc = () => (number | Generator<(arg: number) => void, any, void>)
+>FuncOrGeneratorFunc : () => number | Generator<(arg: number) => void, any, void>
+>arg : number
+
+const f: FuncOrGeneratorFunc = function*() {
+>f : FuncOrGeneratorFunc
+>function*() {  yield (num) => console.log(num); // `num` should be inferred to have type `number`.} : () => Generator<(num: number) => void, void, any>
+
+  yield (num) => console.log(num); // `num` should be inferred to have type `number`.
+>yield (num) => console.log(num) : any
+>(num) => console.log(num) : (num: number) => void
+>num : number
+>console.log(num) : void
+>console.log : (...data: any[]) => void
+>console : Console
+>log : (...data: any[]) => void
+>num : number
+}

tests/cases/compiler/contextualTypeOnYield.ts

@@ -0,0 +1,8 @@
+// @strict: true
+// @target: ES6
+
+type FuncOrGeneratorFunc = () => (number | Generator<(arg: number) => void, any, void>)
+
+const f: FuncOrGeneratorFunc = function*() {
+  yield (num) => console.log(num); // `num` should be inferred to have type `number`.
+}