language-servers/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-09-30 05:45:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Add some basic subscript type inference (#13562)

Browse source 
## Summary

Just for tuples and strings -- the easiest cases. I think most of the
rest require generic support?
This commit is contained in:
Charlie Marsh 2024-09-30 16:50:46 -04:00 • committed by GitHub
parent 32c746bd82
commit c9c748a79e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 227 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

231
crates/red_knot_python_semantic/src/types/infer.rs
View file

@ -1250,6 +1250,42 @@ impl<'db> TypeInferenceBuilder<'db> {
); );
} }
/// Emit a diagnostic declaring that an index is out of bounds for a tuple.
pub(super) fn tuple_index_out_of_bounds_diagnostic(
&mut self,
node: AnyNodeRef,
tuple_ty: Type<'db>,
length: usize,
index: i64,
) {
self.add_diagnostic(
node,
"index-out-of-bounds",
format_args!(
"Index {index} is out of bounds for tuple of type '{}' with length {length}.",
tuple_ty.display(self.db)
),
);
}
/// Emit a diagnostic declaring that an index is out of bounds for a string.
pub(super) fn string_index_out_of_bounds_diagnostic(
&mut self,
node: AnyNodeRef,
string_ty: Type<'db>,
length: usize,
index: i64,
) {
self.add_diagnostic(
node,
"index-out-of-bounds",
format_args!(
"Index {index} is out of bounds for string '{}' with length {length}.",
string_ty.display(self.db)
),
);
}
fn infer_for_statement_definition( fn infer_for_statement_definition(
&mut self, &mut self,
target: &ast::ExprName, target: &ast::ExprName,
@ -2389,11 +2425,127 @@ impl<'db> TypeInferenceBuilder<'db> {
ctx: _, ctx: _,
} = subscript; } = subscript;
self.infer_expression(slice); let value_ty = self.infer_expression(value);
self.infer_expression(value); let slice_ty = self.infer_expression(slice);
// TODO actual subscript support match (value_ty, slice_ty) {
Type::Unknown // Ex) Given `("a", "b", "c", "d")[1]`, return `"b"`
(Type::Tuple(tuple_ty), Type::IntLiteral(int)) if int >= 0 => {
let elements = tuple_ty.elements(self.db);
usize::try_from(int)
.ok()
.and_then(|index| elements.get(index).copied())
.unwrap_or_else(|| {
self.tuple_index_out_of_bounds_diagnostic(
(&**value).into(),
value_ty,
elements.len(),
int,
);
Type::Unknown
})
}
// Ex) Given `("a", "b", "c", "d")[-1]`, return `"c"`
(Type::Tuple(tuple_ty), Type::IntLiteral(int)) if int < 0 => {
let elements = tuple_ty.elements(self.db);
int.checked_neg()
.and_then(|int| usize::try_from(int).ok())
.and_then(|index| elements.len().checked_sub(index))
.and_then(|index| elements.get(index).copied())
.unwrap_or_else(|| {
self.tuple_index_out_of_bounds_diagnostic(
(&**value).into(),
value_ty,
elements.len(),
int,
);
Type::Unknown
})
}
// Ex) Given `("a", "b", "c", "d")[True]`, return `"b"`
(Type::Tuple(tuple_ty), Type::BooleanLiteral(bool)) => {
let elements = tuple_ty.elements(self.db);
let int = i64::from(bool);
elements.get(usize::from(bool)).copied().unwrap_or_else(|| {
self.tuple_index_out_of_bounds_diagnostic(
(&**value).into(),
value_ty,
elements.len(),
int,
);
Type::Unknown
})
}
// Ex) Given `"value"[1]`, return `"a"`
(Type::StringLiteral(literal_ty), Type::IntLiteral(int)) if int >= 0 => {
let literal_value = literal_ty.value(self.db);
usize::try_from(int)
.ok()
.and_then(|index| literal_value.chars().nth(index))
.map(|ch| {
Type::StringLiteral(StringLiteralType::new(
self.db,
ch.to_string().into_boxed_str(),
))
})
.unwrap_or_else(|| {
self.string_index_out_of_bounds_diagnostic(
(&**value).into(),
value_ty,
literal_value.chars().count(),
int,
);
Type::Unknown
})
}
// Ex) Given `"value"[-1]`, return `"e"`
(Type::StringLiteral(literal_ty), Type::IntLiteral(int)) if int < 0 => {
let literal_value = literal_ty.value(self.db);
int.checked_neg()
.and_then(|int| usize::try_from(int).ok())
.and_then(|index| index.checked_sub(1))
.and_then(|index| literal_value.chars().rev().nth(index))
.map(|ch| {
Type::StringLiteral(StringLiteralType::new(
self.db,
ch.to_string().into_boxed_str(),
))
})
.unwrap_or_else(|| {
self.string_index_out_of_bounds_diagnostic(
(&**value).into(),
value_ty,
literal_value.chars().count(),
int,
);
Type::Unknown
})
}
// Ex) Given `"value"[True]`, return `"a"`
(Type::StringLiteral(literal_ty), Type::BooleanLiteral(bool)) => {
let literal_value = literal_ty.value(self.db);
let int = i64::from(bool);
literal_value
.chars()
.nth(usize::from(bool))
.map(|ch| {
Type::StringLiteral(StringLiteralType::new(
self.db,
ch.to_string().into_boxed_str(),
))
})
.unwrap_or_else(|| {
self.string_index_out_of_bounds_diagnostic(
(&**value).into(),
value_ty,
literal_value.chars().count(),
int,
);
Type::Unknown
})
}
_ => Type::Unknown,
}
} }
fn infer_slice_expression(&mut self, slice: &ast::ExprSlice) -> Type<'db> { fn infer_slice_expression(&mut self, slice: &ast::ExprSlice) -> Type<'db> {
@ -6425,6 +6577,77 @@ mod tests {
Ok(()) Ok(())
} }
#[test]
fn subscript_tuple() -> anyhow::Result<()> {
let mut db = setup_db();
db.write_dedented(
"/src/a.py",
"
t = (1, 'a', 'b')
a = t[0]
b = t[1]
c = t[-1]
d = t[-2]
e = t[4]
f = t[-4]
",
)?;
assert_public_ty(&db, "/src/a.py", "a", "Literal[1]");
assert_public_ty(&db, "/src/a.py", "b", "Literal[\"a\"]");
assert_public_ty(&db, "/src/a.py", "c", "Literal[\"b\"]");
assert_public_ty(&db, "/src/a.py", "d", "Literal[\"a\"]");
assert_public_ty(&db, "/src/a.py", "e", "Unknown");
assert_public_ty(&db, "/src/a.py", "f", "Unknown");
assert_file_diagnostics(
&db,
"src/a.py",
&["Index 4 is out of bounds for tuple of type 'tuple[Literal[1], Literal[\"a\"], Literal[\"b\"]]' with length 3.", "Index -4 is out of bounds for tuple of type 'tuple[Literal[1], Literal[\"a\"], Literal[\"b\"]]' with length 3."],
);
Ok(())
}
#[test]
fn subscript_literal_string() -> anyhow::Result<()> {
let mut db = setup_db();
db.write_dedented(
"/src/a.py",
"
s = 'abcde'
a = s[0]
b = s[1]
c = s[-1]
d = s[-2]
e = s[8]
f = s[-8]
",
)?;
assert_public_ty(&db, "/src/a.py", "a", "Literal[\"a\"]");
assert_public_ty(&db, "/src/a.py", "b", "Literal[\"b\"]");
assert_public_ty(&db, "/src/a.py", "c", "Literal[\"e\"]");
assert_public_ty(&db, "/src/a.py", "d", "Literal[\"d\"]");
assert_public_ty(&db, "/src/a.py", "e", "Unknown");
assert_public_ty(&db, "/src/a.py", "f", "Unknown");
assert_file_diagnostics(
&db,
"src/a.py",
&[
"Index 8 is out of bounds for string 'Literal[\"abcde\"]' with length 5.",
"Index -8 is out of bounds for string 'Literal[\"abcde\"]' with length 5.",
],
);
Ok(())
}
#[test] #[test]
fn boolean_or_expression() -> anyhow::Result<()> { fn boolean_or_expression() -> anyhow::Result<()> {
let mut db = setup_db(); let mut db = setup_db();