From 6f52d573ef9fc231076195df58fe298cce4dd9d2 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 28 Oct 2024 10:13:01 -0400 Subject: [PATCH] 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 --- .../resources/mdtest/assignment/annotations.md | 17 +++++++++++++---- .../red_knot_python_semantic/src/types/infer.rs | 16 +++++++++++----- crates/ruff_benchmark/benches/red_knot.rs | 8 +++++++- 3 files changed, 31 insertions(+), 10 deletions(-) diff --git a/crates/red_knot_python_semantic/resources/mdtest/assignment/annotations.md b/crates/red_knot_python_semantic/resources/mdtest/assignment/annotations.md index 923fca96ac..68481fc2a5 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/assignment/annotations.md +++ b/crates/red_knot_python_semantic/resources/mdtest/assignment/annotations.md @@ -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 ``` diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index e9e1f99a43..6c86b336a9 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -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 + } } } diff --git a/crates/ruff_benchmark/benches/red_knot.rs b/crates/ruff_benchmark/benches/red_knot.rs index 657b278079..795c333d7f 100644 --- a/crates/ruff_benchmark/benches/red_knot.rs +++ b/crates/ruff_benchmark/benches/red_knot.rs @@ -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 {