From b7cef6c999108f3139a6a1129d32056b3b710af8 Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Tue, 10 Sep 2024 23:24:19 +0530 Subject: [PATCH] [red-knot] Add heterogeneous tuple type variant (#13295) ## Summary This PR adds a new `Type` variant called `TupleType` which is used for heterogeneous elements. ### Display notes * For an empty tuple, I'm using `tuple[()]` as described in the docs: https://docs.python.org/3/library/typing.html#annotating-tuples * For nested elements, it'll use the literal type instead of builtin type unlike Pyright which does `tuple[Literal[1], tuple[int, int]]` instead of `tuple[Literal[1], tuple[Literal[2], Literal[3]]]`. Also, mypy would give `tuple[builtins.int, builtins.int]` instead of `tuple[Literal[1], Literal[2]]` ## Test Plan Update test case to account for the display change and add cases for multiple elements and nested tuple elements. --------- Co-authored-by: Alex Waygood Co-authored-by: Carl Meyer --- crates/red_knot_python_semantic/src/types.rs | 14 ++++++ .../src/types/display.rs | 17 +++++++ .../src/types/infer.rs | 47 +++++++++++++++---- 3 files changed, 69 insertions(+), 9 deletions(-) diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs index 252b9125b7708..c5f38eb642082 100644 --- a/crates/red_knot_python_semantic/src/types.rs +++ b/crates/red_knot_python_semantic/src/types.rs @@ -195,6 +195,9 @@ pub enum Type<'db> { LiteralString, /// A bytes literal BytesLiteral(BytesLiteralType<'db>), + /// A heterogeneous tuple type, with elements of the given types in source order. + // TODO: Support variable length homogeneous tuple type like `tuple[int, ...]`. + Tuple(TupleType<'db>), // TODO protocols, callable types, overloads, generics, type vars } @@ -362,6 +365,10 @@ impl<'db> Type<'db> { // TODO defer to Type::Instance().member Type::Unknown } + Type::Tuple(_) => { + // TODO: implement tuple methods + Type::Unknown + } } } @@ -473,6 +480,7 @@ impl<'db> Type<'db> { Type::Unknown => Type::Unknown, // TODO intersections Type::Intersection(_) => Type::Unknown, + Type::Tuple(_) => builtins_symbol_ty(db, "tuple"), } } } @@ -658,3 +666,9 @@ pub struct BytesLiteralType<'db> { #[return_ref] value: Box<[u8]>, } + +#[salsa::interned] +pub struct TupleType<'db> { + #[return_ref] + elements: Box<[Type<'db>]>, +} diff --git a/crates/red_knot_python_semantic/src/types/display.rs b/crates/red_knot_python_semantic/src/types/display.rs index 49241154994a6..0d7e2ecf511ce 100644 --- a/crates/red_knot_python_semantic/src/types/display.rs +++ b/crates/red_knot_python_semantic/src/types/display.rs @@ -86,6 +86,23 @@ impl std::fmt::Display for DisplayRepresentation<'_> { escape.bytes_repr().write(f) } + Type::Tuple(tuple) => { + f.write_str("tuple[")?; + let elements = tuple.elements(self.db); + if elements.is_empty() { + f.write_str("()")?; + } else { + let mut first = true; + for element in &**elements { + if !first { + f.write_str(", ")?; + } + first = false; + element.display(self.db).fmt(f)?; + } + } + f.write_str("]") + } } } } diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 02b0efb3b51ec..051e8db2bf627 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -49,7 +49,7 @@ use crate::stdlib::builtins_module_scope; use crate::types::diagnostic::{TypeCheckDiagnostic, TypeCheckDiagnostics}; use crate::types::{ builtins_symbol_ty, definitions_ty, global_symbol_ty, symbol_ty, BytesLiteralType, ClassType, - FunctionType, StringLiteralType, Type, UnionBuilder, + FunctionType, StringLiteralType, TupleType, Type, UnionBuilder, }; use crate::Db; @@ -1553,12 +1553,12 @@ impl<'db> TypeInferenceBuilder<'db> { parenthesized: _, } = tuple; - for elt in elts { - self.infer_expression(elt); - } + let element_types = elts + .iter() + .map(|elt| self.infer_expression(elt)) + .collect::>(); - // TODO generic - builtins_symbol_ty(self.db, "tuple").to_instance() + Type::Tuple(TupleType::new(self.db, element_types.into_boxed_slice())) } fn infer_list_expression(&mut self, list: &ast::ExprList) -> Type<'db> { @@ -4012,7 +4012,7 @@ mod tests { } #[test] - fn tuple_literal() -> anyhow::Result<()> { + fn empty_tuple_literal() -> anyhow::Result<()> { let mut db = setup_db(); db.write_dedented( @@ -4022,8 +4022,37 @@ mod tests { ", )?; - // TODO should be a generic type - assert_public_ty(&db, "/src/a.py", "x", "tuple"); + assert_public_ty(&db, "/src/a.py", "x", "tuple[()]"); + + Ok(()) + } + + #[test] + fn tuple_heterogeneous_literal() -> anyhow::Result<()> { + let mut db = setup_db(); + + db.write_dedented( + "/src/a.py", + " + x = (1, 'a') + y = (1, (2, 3)) + z = (x, 2) + ", + )?; + + assert_public_ty(&db, "/src/a.py", "x", r#"tuple[Literal[1], Literal["a"]]"#); + assert_public_ty( + &db, + "/src/a.py", + "y", + "tuple[Literal[1], tuple[Literal[2], Literal[3]]]", + ); + assert_public_ty( + &db, + "/src/a.py", + "z", + r#"tuple[tuple[Literal[1], Literal["a"]], Literal[2]]"#, + ); Ok(()) }