Skip to content

Commit 2736b6a

Browse files
GlyphackGlyphack
authored
Merge 8816b4f into 6c8101a 
2 parents 6c8101a + 8816b4f commit 2736b6a

8 files changed

Lines changed: 264 additions & 7 deletions

‎crates/ty_python_semantic/resources/mdtest/expression/yield_and_yield_from.md‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ def outer_generator():
9595
A dict literal that is structurally compatible with a `TypedDict` should be accepted.
9696

9797
```py
98-
from typing import Iterator, TypedDict
98+
from typing import Iterator, Generator, TypedDict
9999

100100
class Person(TypedDict):
101101
name: str
@@ -121,6 +121,19 @@ def persons() -> Iterator[Person]:
121121
yield from [{"name": 42}]
122122
```
123123

124+
This also works for return values:
125+
126+
```py
127+
def persons(f: bool) -> Generator[None, None, Person]:
128+
yield
129+
if f:
130+
return {"name": "Bob"}
131+
else:
132+
# error: [invalid-return-type]
133+
# error: [invalid-argument-type]
134+
return {"name": 42}
135+
```
136+
124137
## `yield` expression send type inference
125138

126139
```py
@@ -170,6 +183,7 @@ async def async_iterator_send_none() -> AsyncIterator[int]:
170183

171184
def iterator_yield_from() -> Generator[int, None, int]:
172185
yield from iterator_send_none()
186+
return 1
173187
```
174188

175189
## Error cases
@@ -214,14 +228,14 @@ def sync_returns_async_generator() -> AsyncGenerator[int, str]: # error: [inval
214228
```py
215229
from typing import Generator
216230

217-
# TODO: should emit an error (does not return `str`)
231+
# error: [invalid-return-type]
218232
def invalid_generator1() -> Generator[int, None, str]:
219233
yield 1
220234

221-
# TODO: should emit an error (does not return `int`)
222235
def invalid_generator2() -> Generator[int, None, None]:
223236
yield 1
224237

238+
# error: [invalid-return-type]
225239
return "done"
226240
```
227241

@@ -252,3 +266,17 @@ def outer() -> Generator[int, str, None]:
252266
# error: [invalid-yield] "Send type `int` does not match annotated send type `str`"
253267
yield from inner()
254268
```
269+
270+
### Non generator function with `Generator` annotation
271+
272+
<!-- snapshot-diagnostics -->
273+
274+
```py
275+
from typing import Generator
276+
277+
def non_gen() -> Generator[int, int, None]:
278+
# error: [invalid-return-type]
279+
return 1
280+
281+
reveal_type(non_gen) # revealed: def non_gen() -> Generator[int, int, None]
282+
```

‎crates/ty_python_semantic/resources/mdtest/function/return_type.md‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -533,6 +533,38 @@ def i2() -> typing.Generator:
533533

534534
def j() -> str: # error: [invalid-return-type]
535535
yield 42
536+
537+
def invalid_return_type() -> typing.Generator[None, None, None]:
538+
yield
539+
return "" # error: [invalid-return-type]
540+
```
541+
542+
The return value of the function must be assignable to the return type of the `Generator`. This is
543+
specified in the third type parameter.
544+
545+
```py
546+
def wrong_return() -> typing.Generator[int, int, int]:
547+
yield 1
548+
return "" # error: [invalid-return-type]
549+
```
550+
551+
If the function has no return and it's implicitly returning it is still type checked.
552+
553+
```py
554+
def bare_return_ok() -> typing.Generator[int, int, None]:
555+
yield 1
556+
557+
def missing_return() -> typing.Generator[int, int, int]: # error: [invalid-return-type]
558+
yield 1
559+
```
560+
561+
Iterators must not return anything.
562+
563+
```py
564+
def iterator_must_not_return() -> typing.Iterator[int]:
565+
yield 2
566+
# error: [invalid-return-type]
567+
return "foo"
536568
```
537569

538570
### Asynchronous
@@ -560,6 +592,10 @@ async def i() -> typing.AsyncIterable:
560592

561593
async def j() -> str: # error: [invalid-return-type]
562594
yield 42
595+
596+
async def k() -> typing.AsyncGenerator:
597+
yield 42
598+
return 2 # error: [invalid-syntax] "`return` with value in async generator"
563599
```
564600

565601
## Diagnostics for `empty-body` on non-protocol subclasses of protocol classes

‎crates/ty_python_semantic/resources/mdtest/snapshots/return_type.md_-_Function_return_type_-_Generator_functions_-_Asynchronous_(408134055c24a538).snap‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,10 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/function/return_type.md
3030
15 |
3131
16 | async def j() -> str: # error: [invalid-return-type]
3232
17 | yield 42
33+
18 |
34+
19 | async def k() -> typing.AsyncGenerator:
35+
20 | yield 42
36+
21 | return 2 # error: [invalid-syntax] "`return` with value in async generator"
3337
```
3438

3539
# Diagnostics
@@ -49,3 +53,15 @@ info: See https://docs.python.org/3/glossary.html#term-asynchronous-generator fo
4953
info: rule `invalid-return-type` is enabled by default
5054

5155
```
56+
57+
```
58+
error[invalid-syntax]: `return` with value in async generator
59+
--> src/mdtest_snippet.py:21:5
60+
|
61+
19 | async def k() -> typing.AsyncGenerator:
62+
20 | yield 42
63+
21 | return 2 # error: [invalid-syntax] "`return` with value in async generator"
64+
| ^^^^^^^^
65+
|
66+
67+
```

‎crates/ty_python_semantic/resources/mdtest/snapshots/return_type.md_-_Function_return_type_-_Generator_functions_-_Synchronous_(6a32ec69d15117b8).snap‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,22 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/function/return_type.md
3333
18 |
3434
19 | def j() -> str: # error: [invalid-return-type]
3535
20 | yield 42
36+
21 |
37+
22 | def invalid_return_type() -> typing.Generator[None, None, None]:
38+
23 | yield
39+
24 | return "" # error: [invalid-return-type]
40+
25 | def wrong_return() -> typing.Generator[int, int, int]:
41+
26 | yield 1
42+
27 | return "" # error: [invalid-return-type]
43+
28 | def bare_return_ok() -> typing.Generator[int, int, None]:
44+
29 | yield 1
45+
30 |
46+
31 | def missing_return() -> typing.Generator[int, int, int]: # error: [invalid-return-type]
47+
32 | yield 1
48+
33 | def iterator_must_not_return() -> typing.Iterator[int]:
49+
34 | yield 2
50+
35 | # error: [invalid-return-type]
51+
36 | return "foo"
3652
```
3753

3854
# Diagnostics
@@ -52,3 +68,72 @@ info: See https://docs.python.org/3/glossary.html#term-generator for more detail
5268
info: rule `invalid-return-type` is enabled by default
5369

5470
```
71+
72+
```
73+
error[invalid-return-type]: Return type does not match returned value
74+
--> src/mdtest_snippet.py:22:30
75+
|
76+
20 | yield 42
77+
21 |
78+
22 | def invalid_return_type() -> typing.Generator[None, None, None]:
79+
| ---------------------------------- Expected `None` because of return type
80+
23 | yield
81+
24 | return "" # error: [invalid-return-type]
82+
| ^^ expected `None`, found `Literal[""]`
83+
25 | def wrong_return() -> typing.Generator[int, int, int]:
84+
26 | yield 1
85+
|
86+
info: rule `invalid-return-type` is enabled by default
87+
88+
```
89+
90+
```
91+
error[invalid-return-type]: Return type does not match returned value
92+
--> src/mdtest_snippet.py:25:23
93+
|
94+
23 | yield
95+
24 | return "" # error: [invalid-return-type]
96+
25 | def wrong_return() -> typing.Generator[int, int, int]:
97+
| ------------------------------- Expected `int` because of return type
98+
26 | yield 1
99+
27 | return "" # error: [invalid-return-type]
100+
| ^^ expected `int`, found `Literal[""]`
101+
28 | def bare_return_ok() -> typing.Generator[int, int, None]:
102+
29 | yield 1
103+
|
104+
info: rule `invalid-return-type` is enabled by default
105+
106+
```
107+
108+
```
109+
error[invalid-return-type]: Function always implicitly returns `None`, which is not assignable to return type `int`
110+
--> src/mdtest_snippet.py:31:25
111+
|
112+
29 | yield 1
113+
30 |
114+
31 | def missing_return() -> typing.Generator[int, int, int]: # error: [invalid-return-type]
115+
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
116+
32 | yield 1
117+
33 | def iterator_must_not_return() -> typing.Iterator[int]:
118+
|
119+
info: Consider changing the return annotation to `-> None` or adding a `return` statement
120+
info: rule `invalid-return-type` is enabled by default
121+
122+
```
123+
124+
```
125+
error[invalid-return-type]: Return type does not match returned value
126+
--> src/mdtest_snippet.py:33:35
127+
|
128+
31 | def missing_return() -> typing.Generator[int, int, int]: # error: [invalid-return-type]
129+
32 | yield 1
130+
33 | def iterator_must_not_return() -> typing.Iterator[int]:
131+
| -------------------- Expected `None` because of return type
132+
34 | yield 2
133+
35 | # error: [invalid-return-type]
134+
36 | return "foo"
135+
| ^^^^^ expected `None`, found `Literal["foo"]`
136+
|
137+
info: rule `invalid-return-type` is enabled by default
138+
139+
```

‎crates/ty_python_semantic/resources/mdtest/snapshots/yield_and_yield_from…_-_`yield`_and_`yield_f…_-_Error_cases_-_Non_generator_functi…_(c14a872d57170530).snap‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
---
2+
source: crates/ty_test/src/lib.rs
3+
expression: snapshot
4+
---
5+
6+
---
7+
mdtest name: yield_and_yield_from.md - `yield` and `yield from` - Error cases - Non generator function with `Generator` annotation
8+
mdtest path: crates/ty_python_semantic/resources/mdtest/expression/yield_and_yield_from.md
9+
---
10+
11+
# Python source files
12+
13+
## mdtest_snippet.py
14+
15+
```
16+
1 | from typing import Generator
17+
2 |
18+
3 | def non_gen() -> Generator[int, int, None]:
19+
4 | # error: [invalid-return-type]
20+
5 | return 1
21+
6 |
22+
7 | reveal_type(non_gen) # revealed: def non_gen() -> Generator[int, int, None]
23+
```
24+
25+
# Diagnostics
26+
27+
```
28+
error[invalid-return-type]: Return type does not match returned value
29+
--> src/mdtest_snippet.py:3:18
30+
|
31+
1 | from typing import Generator
32+
2 |
33+
3 | def non_gen() -> Generator[int, int, None]:
34+
| ------------------------- Expected `Generator[int, int, None]` because of return type
35+
4 | # error: [invalid-return-type]
36+
5 | return 1
37+
| ^ expected `Generator[int, int, None]`, found `Literal[1]`
38+
6 |
39+
7 | reveal_type(non_gen) # revealed: def non_gen() -> Generator[int, int, None]
40+
|
41+
info: rule `invalid-return-type` is enabled by default
42+
43+
```

‎crates/ty_python_semantic/src/types.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4823,7 +4823,7 @@ impl<'db> Type<'db> {
48234823
Some(GeneratorTypes {
48244824
yield_ty: Some(*yield_ty),
48254825
send_ty: Some(Type::none(db)),
4826-
return_ty: Some(Type::unknown()),
4826+
return_ty: Some(Type::none(db)),
48274827
})
48284828
} else {
48294829
None

‎crates/ty_python_semantic/src/types/infer/builder.rs‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4601,9 +4601,21 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
46014601
// the expected type passed should be the "raw" type,
46024602
// i.e. type variables in the return type are non-inferable,
46034603
// and the return types of async functions are not wrapped in `CoroutineType[...]`.
4604-
TypeContext::new(Some(
4605-
func.last_definition_raw_signature(self.db()).return_ty,
4606-
))
4604+
let return_ty = func.last_definition_raw_signature(self.db()).return_ty;
4605+
4606+
// For generator functions, the declared return type is e.g.
4607+
// `Generator[YieldType, SendType, ReturnType]`. The type context
4608+
// for a `return` statement should be the `ReturnType` type parameter
4609+
let file_scope_id = self.scope().file_scope_id(self.db());
4610+
let context_ty = if file_scope_id.is_generator_function(self.index) {
4611+
return_ty
4612+
.generator_return_type(self.db())
4613+
.unwrap_or(return_ty)
4614+
} else {
4615+
return_ty
4616+
};
4617+
4618+
TypeContext::new(Some(context_ty))
46074619
})
46084620
.unwrap_or_default()
46094621
} else {

‎crates/ty_python_semantic/src/types/infer/builder/function.rs‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,43 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
114114
declared_ty,
115115
);
116116
}
117+
118+
if let Some(expected_return_ty) = declared_ty.generator_return_type(db) {
119+
for invalid in
120+
self.return_types_and_ranges
121+
.iter()
122+
.copied()
123+
.filter(|actual_return_ty| {
124+
!actual_return_ty.ty.is_assignable_to(db, expected_return_ty)
125+
})
126+
{
127+
report_invalid_return_type(
128+
&self.context,
129+
invalid.range,
130+
returns.range(),
131+
expected_return_ty,
132+
invalid.ty,
133+
);
134+
}
135+
136+
if self
137+
.index
138+
.use_def_map(scope_id)
139+
.can_implicitly_return_none(db)
140+
&& !Type::none(db).is_assignable_to(db, expected_return_ty)
141+
{
142+
let no_return = self.return_types_and_ranges.is_empty();
143+
report_implicit_return_type(
144+
&self.context,
145+
returns.range(),
146+
expected_return_ty,
147+
false,
148+
None,
149+
no_return,
150+
);
151+
}
152+
}
153+
117154
return;
118155
}
119156

0 commit comments

Comments
 (0)