Support inference for PEP 604 union annotations (#13964)

## Summary

Supports return type inference for, e.g., `def f() -> int | None:`.

---------

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
This commit is contained in:
Charlie Marsh 2024-10-28 10:13:01 -04:00 committed by GitHub
parent c593ccb529
commit 6f52d573ef
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 31 additions and 10 deletions

View file

@ -23,12 +23,21 @@ x: int
x = "foo" # error: [invalid-assignment] "Object of type `Literal["foo"]` is not assignable to `int`"
```
## PEP-604 annotations not yet supported
## PEP-604 annotations are supported
```py
def f() -> str | None:
def foo() -> str | int | None:
return None
# TODO: should be `str | None` (but Todo is better than `Unknown`)
reveal_type(f()) # revealed: @Todo
reveal_type(foo()) # revealed: str | int | None
def bar() -> str | str | None:
return None
reveal_type(bar()) # revealed: str | None
def baz() -> str | str:
return "Hello, world!"
reveal_type(baz()) # revealed: str
```

View file

@ -3494,14 +3494,20 @@ impl<'db> TypeInferenceBuilder<'db> {
Type::Todo
}
// TODO PEP-604 unions
ast::Expr::BinOp(binary) => {
self.infer_binary_expression(binary);
#[allow(clippy::single_match_else)]
match binary.op {
// PEP-604 unions are okay
ast::Operator::BitOr => Type::Todo,
// PEP-604 unions are okay, e.g., `int | str`
ast::Operator::BitOr => {
let left_ty = self.infer_type_expression(&binary.left);
let right_ty = self.infer_type_expression(&binary.right);
UnionType::from_elements(self.db, [left_ty, right_ty])
}
// anything else is an invalid annotation:
_ => Type::Unknown,
_ => {
self.infer_binary_expression(binary);
Type::Unknown
}
}
}

View file

@ -32,11 +32,16 @@ static EXPECTED_DIAGNOSTICS: &[&str] = &[
"/src/tomllib/_parser.py:98:12: Name `char` used when possibly not defined",
"/src/tomllib/_parser.py:101:12: Name `char` used when possibly not defined",
"/src/tomllib/_parser.py:104:14: Name `char` used when possibly not defined",
"/src/tomllib/_parser.py:108:17: Conflicting declared types for `second_char`: Unknown, str | None",
"/src/tomllib/_parser.py:115:14: Name `char` used when possibly not defined",
"/src/tomllib/_parser.py:126:12: Name `char` used when possibly not defined",
"/src/tomllib/_parser.py:267:9: Conflicting declared types for `char`: Unknown, str | None",
"/src/tomllib/_parser.py:348:20: Name `nest` used when possibly not defined",
"/src/tomllib/_parser.py:353:5: Name `nest` used when possibly not defined",
"/src/tomllib/_parser.py:353:5: Method `__getitem__` of type `Unbound | @Todo` is not callable on object of type `Unbound | @Todo`",
"/src/tomllib/_parser.py:364:9: Conflicting declared types for `char`: Unknown, str | None",
"/src/tomllib/_parser.py:381:13: Conflicting declared types for `char`: Unknown, str | None",
"/src/tomllib/_parser.py:395:9: Conflicting declared types for `char`: Unknown, str | None",
"/src/tomllib/_parser.py:453:24: Name `nest` used when possibly not defined",
"/src/tomllib/_parser.py:455:9: Name `nest` used when possibly not defined",
"/src/tomllib/_parser.py:455:9: Method `__getitem__` of type `Unbound | @Todo` is not callable on object of type `Unbound | @Todo`",
@ -45,7 +50,8 @@ static EXPECTED_DIAGNOSTICS: &[&str] = &[
"/src/tomllib/_parser.py:573:12: Name `char` used when possibly not defined",
"/src/tomllib/_parser.py:579:12: Name `char` used when possibly not defined",
"/src/tomllib/_parser.py:580:63: Name `char` used when possibly not defined",
"/src/tomllib/_parser.py:629:38: Name `datetime_obj` used when possibly not defined"
"/src/tomllib/_parser.py:590:9: Conflicting declared types for `char`: Unknown, str | None",
"/src/tomllib/_parser.py:629:38: Name `datetime_obj` used when possibly not defined",
];
fn get_test_file(name: &str) -> TestFile {