Skip to content

Commit

Permalink
[red-knot] Add heterogeneous tuple type variant (#13295)
Browse files Browse the repository at this point in the history
## 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 <Alex.Waygood@Gmail.com>
Co-authored-by: Carl Meyer <carl@astral.sh>
  • Loading branch information
3 people authored Sep 10, 2024
1 parent 110193a commit b7cef6c
Show file tree
Hide file tree
Showing 3 changed files with 69 additions and 9 deletions.
14 changes: 14 additions & 0 deletions crates/red_knot_python_semantic/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down Expand Up @@ -362,6 +365,10 @@ impl<'db> Type<'db> {
// TODO defer to Type::Instance(<bytes from typeshed>).member
Type::Unknown
}
Type::Tuple(_) => {
// TODO: implement tuple methods
Type::Unknown
}
}
}

Expand Down Expand Up @@ -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"),
}
}
}
Expand Down Expand Up @@ -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>]>,
}
17 changes: 17 additions & 0 deletions crates/red_knot_python_semantic/src/types/display.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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("]")
}
}
}
}
Expand Down
47 changes: 38 additions & 9 deletions crates/red_knot_python_semantic/src/types/infer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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::<Vec<_>>();

// 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> {
Expand Down Expand Up @@ -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(
Expand All @@ -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(())
}
Expand Down

0 comments on commit b7cef6c

Please sign in to comment.