[ty] Check return type of generator functions

[Glyphack]

Summary

Perform return type checking for generator functions.

part of astral-sh/ty#1718

Test Plan

The two newly added diagnostics in ecosystem result are correct the rest were marked as flaky. So looks good.
I updated the tests for return_types.md.

[astral-sh-bot]

Typing conformance results improved 🎉

The percentage of diagnostics emitted that were expected errors increased from 85.31% to 85.34%. The percentage of expected errors that received a diagnostic increased from 78.79% to 78.98%. The number of fully passing files improved from 65/133 to 66/133.

Summary

How are test cases classified?

Each test case represents one expected error annotation or a group of annotations sharing a tag. Counts are per test case, not per diagnostic — multiple diagnostics on the same line count as one. Required annotations (E) are true positives when ty flags the expected location and false negatives when it does not. Optional annotations (E?) are true positives when flagged but true negatives (not false negatives) when not. Tagged annotations (E[tag]) require ty to flag exactly one of the tagged lines; tagged multi-annotations (E[tag+]) allow any number up to the tag count. Flagging unexpected locations counts as a false positive.

Metric	Old	New	Diff	Outcome
True Positives	836	838	+2	⏫ (✅)
False Positives	144	144	+0
False Negatives	225	223	-2	⏬ (✅)
Total Diagnostics	1057	1059	+2	⏫
Precision	85.31%	85.34%	+0.03%	⏫ (✅)
Recall	78.79%	78.98%	+0.19%	⏫ (✅)
Passing Files	65/133	66/133	+1	⏫ (✅)

Test file breakdown

1 file altered

File	True Positives	False Positives	False Negatives	Status
annotations_generators.py	10 (+2) ✅	0	0 (-2) ✅	✅ Newly Passing 🎉
Total (all files)	838 (+2) ✅	144	223 (-2) ✅	66/133

True positives added (2)

2 diagnostics

Test case	Diff
annotations_generators.py:51	+error[invalid-return-type] Function can implicitly return `None`, which is not assignable to return type `C`
annotations_generators.py:54	+error[invalid-return-type] Return type does not match returned value: expected `C`, found `Literal[False]`

[astral-sh-bot]

Memory usage report

Memory usage unchanged ✅

[astral-sh-bot]

ecosystem-analyzer results

Lint rule	Added	Removed	Changed
invalid-await	0	40	0
invalid-return-type	2	1	0
Total	2	41	0

Changes in flaky projects detected. Raw diff output excludes flaky projects; see the HTML report for details.

Full report with detailed diff (timing results)

[sharkdp]

Thank you!

I think there's potentially a similar problem here like we had in the previous PR. Can we add a test where the generator function is annotated with Generator[…, …, SomeTypedDict] and then try to return a dict literal that is compatible with that TypedDict (see similar test case in yield PR)? We want to make sure that we infer the expression after the return keyword with the type context of the generator return type, not the generator type.

[Glyphack]

The problem exists here, sorry I didn't apply the same thing with last PR. I fixed it by adding this logic to infer_return_statement.

[sharkdp]

Thank you. This looks good, but I think we should think about and test a few more cases here (see inline comments).

[sharkdp]

It looks like we should address this TODO comment in this PR?

[Glyphack]

I think parts of it are already fixed. And I got a question about behavior of it.

we should iterate over the yield expressions and return statements in the function to check that they are consistent with the type argument provided.

Check yield values was done in #23796
Check return values is done in this PR.

Unless we want to record type when we see yield expression then perform the check here. Similar to what we do with return_types_and_ranges.

Once we do this, the .to_instance_unknown call below should be replaced with .to_specialized_instance.

I assume the logic here would be union all inferred return types and use that for return_ty and then union all inferred send types and use it for send_ty and here specialize the instance to Generator[union_of_yields, Unknown, union_of_returns].

And this require to collect the yielded types. So I guess it makes it easier to do both in one go.

---

If we use specialize the instance here then the function annotation is checked against the inferred type.
But doing this would cause some behavior that I'm not sure about:

Here we would emit following diagnostics:

from typing import Generator
# error invalid return because inferred type is `Generator[None, None, int]` but the annotated type is `Generator[None, None, str]` (1)
def a(b: bool) -> Generator[None, None, str]:
    yield
    if b:
		# error invalid return type because the annotation is str
        return 1
    else:
		# error invalid return type because the annotation is str
        return 0

Pyright Playground
Pyrefly Sandbox

(1) we don't do this right now but I expect to do once the TODO is implemented. The reason we don't give this error even with specialization is this:

inferred return GeneratorType[None, None, int]
expected Generator[None, None, str]
inferred assignable to expected: true

Should we emit both errors here? Both saying the returned values should be str and also saying the function annotation should specify int for return_ty type parameter?

But I feel this diagnostic is more useful when the function is annotated with something that cannot be assigned to generator type. In that case we don't emit both diagnostic on annotation and each return statement.

[sharkdp]

I agree that there are some open design questions here.

I assume the logic here would be union all inferred return types and use that for return_ty and then union all inferred send types and use it for send_ty and here specialize the instance to Generator[union_of_yields, Unknown, union_of_returns].

I think so, yes.

If we use specialize the instance here then the function annotation is checked against the inferred type.
But doing this would cause some behavior that I'm not sure about: […]
Should we emit both errors here?

No, definitely not. If there is a Generator[…] annotation, we should only emit errors on the return statements.

But I feel this diagnostic is more useful when the function is annotated with something that cannot be assigned to generator type.

Right. And doing this type inference would be valuable once we add return type inference.

In any case, it seems fine to postpone this for now. Thanks for thinking about it and writing it all down.

[Glyphack]

I was planning to go over this and the other todo after the PRs. I subscribed to type inference PR once that's merged I'll check the approach and do the same for generator functions and resolve this comment.