Skip to content

Commit

Permalink
[red-knot] Add heterogeneous tuple type variant
Browse files Browse the repository at this point in the history
  • Loading branch information
dhruvmanila committed Sep 9, 2024
1 parent 62c7d8f commit 2121346
Show file tree
Hide file tree
Showing 3 changed files with 68 additions and 7 deletions.
13 changes: 13 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,8 @@ pub enum Type<'db> {
LiteralString,
/// A bytes literal
BytesLiteral(BytesLiteralType<'db>),
/// A heterogeneous tuple type, with elements of the given types in source order.
Tuple(TupleType<'db>),
// TODO protocols, callable types, overloads, generics, type vars
}

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

Expand Down Expand Up @@ -473,6 +479,7 @@ impl<'db> Type<'db> {
Type::Unknown => Type::Unknown,
// TODO intersections
Type::Intersection(_) => Type::Unknown,
Type::Tuple(_) => typeshed_symbol_ty(db, "tuple"),
}
}
}
Expand Down Expand Up @@ -658,3 +665,9 @@ pub struct BytesLiteralType<'db> {
#[return_ref]
value: Box<[u8]>,
}

#[salsa::interned]
pub struct TupleType<'db> {
#[return_ref]
elements: Vec<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 tuple.elements(self.db) {
if !first {
f.write_str(", ")?;
}
first = false;
element.display(self.db).fmt(f)?;
}
}
f.write_str("]")
}
}
}
}
Expand Down
45 changes: 38 additions & 7 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,13 @@ 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))
}

fn infer_list_expression(&mut self, list: &ast::ExprList) -> Type<'db> {
Expand Down Expand Up @@ -4012,7 +4013,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 @@ -4023,7 +4024,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 2121346

Please sign in to comment.