Inconsistent Type Inference with yield from Expressions

[rayansostenes]

I’ve encountered an issue with Pyright's handling of type inference when using yield from. Below are a few examples that illustrate the problem.

Normal Yield Expression

def func_a():
    for i in (1, 2, 3):
        yield i

reveal_type(func_a())
#           ^^^^^^^^
#           Type of "func_a()" is "Generator[Literal[1, 2, 3], Any, None]"

In this case, Pyright correctly infers the type of func_a() as Generator[Literal[1, 2, 3], Any, None].

Yield from Expression

def func_b():
#   ^^^^^^
#   ERROR: Return type, "Generator[Literal[1, 2, 3], Unknown, None]", is partially unknown
    yield from (1, 2, 3)

reveal_type(func_b())
#           ^^^^^^^^
#           Type of "func_b()" is "Generator[Literal[1, 2, 3], Unknown, None]"

Here, Pyright is unable to fully infer the return type of func_b(), producing an error due to the partially unknown type Generator[Literal[1, 2, 3], Unknown, None].

Fixing the Error with Explicit Type Hint

I can fix the issue by explicitly specifying the return type of func_b() as follows:

from collections.abc import Generator
from typing import Never

def func_b() -> Generator[int, None, None]:
    yield from (1, 2, 3)

However, I’m trying to understand why Pyright is unable to infer the return type of func_b() without the explicit type hint, even though the code is conceptually similar to the first example.

yield from Inference Issue

Moreover, Pyright seems unable to infer the return type of a yield from expression, as shown in the following example:

def func_c():
    x = yield from (1, 2, 3)
    reveal_type(x)
    #           ^
    #           Type of "x" is "Unknown"

In this case, x is inferred as Unknown.

Correct Type Inference for Yielded Values, But Incomplete for send

Interestingly, if I pass a generator with an explicit type hint, Pyright correctly infers the return type of the function, but it still can’t infer the send type:

def func_d(gen: Generator[int, str, str]):
#   ^^^^^^
#   ERROR: Return type, "Generator[int, Unknown, None]", is partially unknown
    y = yield from gen
    reveal_type(y)
    #           ^
    #           Type of "y" is "str"

And if I deliberately mistype the send value, Pyright correctly catches the error:

def func_e(gen: Generator[int, str, str]) -> Generator[int, int, str]:
    result = yield from gen
            #  ^^^
            #  ERROR:
            #  Return type of generator function must be compatible with "Generator[int, Any, Any]"
            #    "Generator[int, str, Unknown]" is not assignable to "Generator[int, int, str]"
            #      Type parameter "_SendT_contra@Generator" is contravariant, but "str" is not a supertype of "int"
            #        "int" is not assignable to "str"
    return result

Question

Am I missing something, or is this a limitation in Pyright’s type inference? Should Pyright be able to infer the return type in these yield from cases, or is this a known bug?

[erictraut]

I don't see how a type checker could infer the generator's send type in this case. It can infer the yield type and the return type based on the function implementation, but the send type depends on the values passed to the send method on the generator object. That's not something a static type checker can know by looking at the implementation of the generator function.

The same limitation exists for a generator with a simple yield statement. Pyright infers a send type of Any rather than Unknown in this case. Arguably, it should infer Unknown for the send type in this case as well, but this was deemed too noisy in strict mode. In the case of a yield from, it seems more reasonable to infer Unknown because the operand for a yield from is normally a fully-typed generator function, not an expression whose type needs to be inferred.

[rayansostenes]

After exploring Pyright's internals and PEP 380, I found an open issue on the CPython repo: [https://github.com/python/cpython/issues/123521]. It seems that the current type stub for generators could be improved. Specifically, the second type argument _SendT_contra should be _SendT_contra | None, as None is always an accepted value when calling .send() on a generator. I made some changes to Pyright's code and was able to get the type inference working for the most common cases.

diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
index 823f891fe..437de9add 100644
--- a/packages/pyright-internal/src/analyzer/typeEvaluator.ts
+++ b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
@@ -19275,6 +19275,7 @@ export function createTypeEvaluator(
                     // Is it a generator?
                     if (functionDecl?.isGenerator) {
                         const inferredYieldTypes: Type[] = [];
+                        const inferredSendTypes: Type[] = [];
                         let useAwaitableGenerator = false;
                         let isYieldResultUsed = false;
 
@@ -19284,6 +19285,18 @@ export function createTypeEvaluator(
                                     if (yieldNode.nodeType === ParseNodeType.YieldFrom) {
                                         isYieldResultUsed = true;
                                         const iteratorTypeResult = getTypeOfExpression(yieldNode.d.expr);
+                                        if (
+                                            isClassInstance(iteratorTypeResult.type) &&
+                                            ClassType.isBuiltIn(iteratorTypeResult.type, 'Generator')
+                                        ) {
+                                            const sendType =
+                                                (iteratorTypeResult.type.priv.typeArgs &&
+                                                    iteratorTypeResult.type.priv.typeArgs.length >= 2 &&
+                                                    iteratorTypeResult.type.priv.typeArgs[1]) ||
+                                                undefined;
+
+                                            inferredSendTypes.push(sendType ?? UnknownType.create());
+                                        }
                                         if (
                                             isClassInstance(iteratorTypeResult.type) &&
                                             ClassType.isBuiltIn(iteratorTypeResult.type, 'Coroutine')
@@ -19341,7 +19354,18 @@ export function createTypeEvaluator(
                             // if the function never uses the value and Unknown if it does.
                             // This eliminates any "partially unknown" errors in strict mode
                             // in the common case.
-                            const sendType = isYieldResultUsed ? UnknownType.create() : AnyType.create();
+                            let sendType: Type;
+                            if (isYieldResultUsed) {
+                                const combinedType =
+                                    inferredSendTypes.length > 0 ? combineTypes(inferredSendTypes) : getNoneType();
+                                if (isUnion(combinedType)) {
+                                    sendType = UnknownType.create();
+                                } else {
+                                    sendType = combinedType;
+                                }
+                            } else {
+                                sendType = AnyType.create();
+                            }
 
                             typeArgs.push(inferredYieldType, sendType, inferredReturnType);
 
diff --git a/packages/pyright-internal/src/tests/samples/generator17.py b/packages/pyright-internal/src/tests/samples/generator17.py
new file mode 100644
index 000000000..c2b32d1af
--- /dev/null
+++ b/packages/pyright-internal/src/tests/samples/generator17.py
@@ -0,0 +1,18 @@
+# pyright: strict
+
+from typing import Generator
+
+
+def gen1() -> Generator[int, str, None]: ...
+
+
+def func1():
+    yield from [1, 2, 3]
+
+
+def func2(gen: Generator[int, str, None]):
+    yield from gen
+
+
+reveal_type(func1(), expected_text="Generator[int, None, None]")
+reveal_type(func2(gen1()), expected_text="Generator[int, str, None]")
diff --git a/packages/pyright-internal/src/tests/samples/generator3.py b/packages/pyright-internal/src/tests/samples/generator3.py
index 98ca70f81..d919343af 100644
--- a/packages/pyright-internal/src/tests/samples/generator3.py
+++ b/packages/pyright-internal/src/tests/samples/generator3.py
@@ -33,7 +33,7 @@ def generator3():
     return 0
 
 
-reveal_type(generator3(), expected_text="Generator[Literal[3], Unknown, Literal[0]]")
+reveal_type(generator3(), expected_text="Generator[Literal[3], None, Literal[0]]")
 
 
 def consumer1() -> ClassB:
diff --git a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts
index 40bc83799..180355a96 100644
--- a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts
+++ b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts
@@ -139,6 +139,12 @@ test('Generator16', () => {
     TestUtils.validateResults(analysisResults, 1);
 });
 
+test('Generator17', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['generator17.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('Await1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['await1.py']);
 

[erictraut]

I'm not convinced this is the correct behavior. As I said above, I don't think a type checker can (or should attempt to) infer the send type. None is not necessarily the correct send type, so Unknown is preferable. You can override this with an explicit return type annotation if you'd like.