[red-knot] Infer target types for unpacked tuple assignment (#13316)

## Summary

This PR adds support for unpacking tuple expression in an assignment
statement where the target expression can be a tuple or a list (the
allowed sequence targets).

The implementation introduces a new `infer_assignment_target` which can
then be used for other targets like the ones in for loops as well. This
delegates it to the `infer_definition`. The final implementation uses a
recursive function that visits the target expression in source order and
compares the variable node that corresponds to the definition. At the
same time, it keeps track of where it is on the assignment value type.

The logic also accounts for the number of elements on both sides such
that it matches even if there's a gap in between. For example, if
there's a starred expression like `(a, *b, c) = (1, 2, 3)`, then the
type of `a` will be `Literal[1]` and the type of `b` will be
`Literal[2]`.

There are a couple of follow-ups that can be done:
* Use this logic for other target positions like `for` loop
* Add diagnostics for mis-match length between LHS and RHS

## Test Plan

Add various test cases using the new markdown test framework.
Validate that existing test cases pass.

---------

Co-authored-by: Carl Meyer <carl@astral.sh>
This commit is contained in:
Dhruv Manilawala 2024-10-16 00:37:11 +05:30 committed by GitHub
parent d77480768d
commit b16f665a81
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 525 additions and 46 deletions

View file

@ -0,0 +1,273 @@
# Unpacking
## Tuple
### Simple tuple
```py
(a, b, c) = (1, 2, 3)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
reveal_type(c) # revealed: Literal[3]
```
### Simple list
```py
[a, b, c] = (1, 2, 3)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
reveal_type(c) # revealed: Literal[3]
```
### Simple mixed
```py
[a, (b, c), d] = (1, (2, 3), 4)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
reveal_type(c) # revealed: Literal[3]
reveal_type(d) # revealed: Literal[4]
```
### Multiple assignment
```py
a, b = c = 1, 2
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
reveal_type(c) # revealed: tuple[Literal[1], Literal[2]]
```
### Nested tuple with unpacking
```py
(a, (b, c), d) = (1, (2, 3), 4)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
reveal_type(c) # revealed: Literal[3]
reveal_type(d) # revealed: Literal[4]
```
### Nested tuple without unpacking
```py
(a, b, c) = (1, (2, 3), 4)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: tuple[Literal[2], Literal[3]]
reveal_type(c) # revealed: Literal[4]
```
### Uneven unpacking (1)
```py
# TODO: Add diagnostic (there aren't enough values to unpack)
(a, b, c) = (1, 2)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
reveal_type(c) # revealed: Unknown
```
### Uneven unpacking (2)
```py
# TODO: Add diagnostic (too many values to unpack)
(a, b) = (1, 2, 3)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
```
### Starred expression (1)
```py
# TODO: Add diagnostic (need more values to unpack)
# TODO: Remove 'not-iterable' diagnostic
[a, *b, c, d] = (1, 2) # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: Literal[1]
# TODO: Should be list[Any] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: Literal[2]
reveal_type(d) # revealed: Unknown
```
### Starred expression (2)
```py
[a, *b, c] = (1, 2) # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: Literal[1]
# TODO: Should be list[Any] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: Literal[2]
```
### Starred expression (3)
```py
# TODO: Remove 'not-iterable' diagnostic
[a, *b, c] = (1, 2, 3) # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: Literal[1]
# TODO: Should be list[int] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: Literal[3]
```
### Starred expression (4)
```py
# TODO: Remove 'not-iterable' diagnostic
[a, *b, c, d] = (1, 2, 3, 4, 5, 6) # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: Literal[1]
# TODO: Should be list[int] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: Literal[5]
reveal_type(d) # revealed: Literal[6]
```
### Starred expression (5)
```py
# TODO: Remove 'not-iterable' diagnostic
[a, b, *c] = (1, 2, 3, 4) # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
# TODO: Should be list[int] once support for assigning to starred expression is added
reveal_type(c) # revealed: @Todo
```
### Non-iterable unpacking
TODO: Remove duplicate diagnostics. This is happening because for a sequence-like
assignment target, multiple definitions are created and the inference engine runs
on each of them which results in duplicate diagnostics.
```py
# error: "Object of type `Literal[1]` is not iterable"
# error: "Object of type `Literal[1]` is not iterable"
a, b = 1
reveal_type(a) # revealed: Unknown
reveal_type(b) # revealed: Unknown
```
### Custom iterator unpacking
```py
class Iterator:
def __next__(self) -> int:
return 42
class Iterable:
def __iter__(self) -> Iterator:
return Iterator()
(a, b) = Iterable()
reveal_type(a) # revealed: int
reveal_type(b) # revealed: int
```
### Custom iterator unpacking nested
```py
class Iterator:
def __next__(self) -> int:
return 42
class Iterable:
def __iter__(self) -> Iterator:
return Iterator()
(a, (b, c), d) = (1, Iterable(), 2)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: int
reveal_type(c) # revealed: int
reveal_type(d) # revealed: Literal[2]
```
## String
### Simple unpacking
```py
a, b = 'ab'
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
```
### Uneven unpacking (1)
```py
# TODO: Add diagnostic (there aren't enough values to unpack)
a, b, c = 'ab'
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
reveal_type(c) # revealed: Unknown
```
### Uneven unpacking (2)
```py
# TODO: Add diagnostic (too many values to unpack)
a, b = 'abc'
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
```
### Starred expression (1)
```py
# TODO: Add diagnostic (need more values to unpack)
# TODO: Remove 'not-iterable' diagnostic
(a, *b, c, d) = "ab" # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: LiteralString
# TODO: Should be list[LiteralString] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: LiteralString
reveal_type(d) # revealed: Unknown
```
### Starred expression (2)
```py
(a, *b, c) = "ab" # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: LiteralString
# TODO: Should be list[Any] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: LiteralString
```
### Starred expression (3)
```py
# TODO: Remove 'not-iterable' diagnostic
(a, *b, c) = "abc" # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: LiteralString
# TODO: Should be list[LiteralString] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: LiteralString
```
### Starred expression (4)
```py
# TODO: Remove 'not-iterable' diagnostic
(a, *b, c, d) = "abcdef" # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: LiteralString
# TODO: Should be list[LiteralString] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: LiteralString
reveal_type(d) # revealed: LiteralString
```
### Starred expression (5)
```py
# TODO: Remove 'not-iterable' diagnostic
(a, b, *c) = "abcd" # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
# TODO: Should be list[int] once support for assigning to starred expression is added
reveal_type(c) # revealed: @Todo
```

View file

@ -994,7 +994,7 @@ class C[T]:
let ast::Expr::NumberLiteral(ast::ExprNumberLiteral { let ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
value: ast::Number::Int(num), value: ast::Number::Int(num),
.. ..
}) = &*assignment.assignment().value }) = assignment.value()
else { else {
panic!("should be a number literal") panic!("should be a number literal")
}; };

View file

@ -28,8 +28,8 @@ use crate::Db;
use super::constraint::{Constraint, PatternConstraint}; use super::constraint::{Constraint, PatternConstraint};
use super::definition::{ use super::definition::{
DefinitionCategory, ExceptHandlerDefinitionNodeRef, MatchPatternDefinitionNodeRef, AssignmentKind, DefinitionCategory, ExceptHandlerDefinitionNodeRef,
WithItemDefinitionNodeRef, MatchPatternDefinitionNodeRef, WithItemDefinitionNodeRef,
}; };
pub(super) struct SemanticIndexBuilder<'db> { pub(super) struct SemanticIndexBuilder<'db> {
@ -566,11 +566,22 @@ where
debug_assert!(self.current_assignment.is_none()); debug_assert!(self.current_assignment.is_none());
self.visit_expr(&node.value); self.visit_expr(&node.value);
self.add_standalone_expression(&node.value); self.add_standalone_expression(&node.value);
self.current_assignment = Some(node.into()); for (target_index, target) in node.targets.iter().enumerate() {
for target in &node.targets { let kind = match target {
ast::Expr::List(_) | ast::Expr::Tuple(_) => Some(AssignmentKind::Sequence),
ast::Expr::Name(_) => Some(AssignmentKind::Name),
_ => None,
};
if let Some(kind) = kind {
self.current_assignment = Some(CurrentAssignment::Assign {
assignment: node,
target_index,
kind,
});
}
self.visit_expr(target); self.visit_expr(target);
self.current_assignment = None;
} }
self.current_assignment = None;
} }
ast::Stmt::AnnAssign(node) => { ast::Stmt::AnnAssign(node) => {
debug_assert!(self.current_assignment.is_none()); debug_assert!(self.current_assignment.is_none());
@ -815,12 +826,18 @@ where
let symbol = self.add_symbol(id.clone()); let symbol = self.add_symbol(id.clone());
if is_definition { if is_definition {
match self.current_assignment { match self.current_assignment {
Some(CurrentAssignment::Assign(assignment)) => { Some(CurrentAssignment::Assign {
assignment,
target_index,
kind,
}) => {
self.add_definition( self.add_definition(
symbol, symbol,
AssignmentDefinitionNodeRef { AssignmentDefinitionNodeRef {
assignment, assignment,
target: name_node, target_index,
name: name_node,
kind,
}, },
); );
} }
@ -1045,7 +1062,11 @@ where
#[derive(Copy, Clone, Debug)] #[derive(Copy, Clone, Debug)]
enum CurrentAssignment<'a> { enum CurrentAssignment<'a> {
Assign(&'a ast::StmtAssign), Assign {
assignment: &'a ast::StmtAssign,
target_index: usize,
kind: AssignmentKind,
},
AnnAssign(&'a ast::StmtAnnAssign), AnnAssign(&'a ast::StmtAnnAssign),
AugAssign(&'a ast::StmtAugAssign), AugAssign(&'a ast::StmtAugAssign),
For(&'a ast::StmtFor), For(&'a ast::StmtFor),
@ -1057,12 +1078,6 @@ enum CurrentAssignment<'a> {
WithItem(&'a ast::WithItem), WithItem(&'a ast::WithItem),
} }
impl<'a> From<&'a ast::StmtAssign> for CurrentAssignment<'a> {
fn from(value: &'a ast::StmtAssign) -> Self {
Self::Assign(value)
}
}
impl<'a> From<&'a ast::StmtAnnAssign> for CurrentAssignment<'a> { impl<'a> From<&'a ast::StmtAnnAssign> for CurrentAssignment<'a> {
fn from(value: &'a ast::StmtAnnAssign) -> Self { fn from(value: &'a ast::StmtAnnAssign) -> Self {
Self::AnnAssign(value) Self::AnnAssign(value)

View file

@ -161,7 +161,9 @@ pub(crate) struct ImportFromDefinitionNodeRef<'a> {
#[derive(Copy, Clone, Debug)] #[derive(Copy, Clone, Debug)]
pub(crate) struct AssignmentDefinitionNodeRef<'a> { pub(crate) struct AssignmentDefinitionNodeRef<'a> {
pub(crate) assignment: &'a ast::StmtAssign, pub(crate) assignment: &'a ast::StmtAssign,
pub(crate) target: &'a ast::ExprName, pub(crate) target_index: usize,
pub(crate) name: &'a ast::ExprName,
pub(crate) kind: AssignmentKind,
} }
#[derive(Copy, Clone, Debug)] #[derive(Copy, Clone, Debug)]
@ -224,12 +226,17 @@ impl DefinitionNodeRef<'_> {
DefinitionNodeRef::NamedExpression(named) => { DefinitionNodeRef::NamedExpression(named) => {
DefinitionKind::NamedExpression(AstNodeRef::new(parsed, named)) DefinitionKind::NamedExpression(AstNodeRef::new(parsed, named))
} }
DefinitionNodeRef::Assignment(AssignmentDefinitionNodeRef { assignment, target }) => { DefinitionNodeRef::Assignment(AssignmentDefinitionNodeRef {
DefinitionKind::Assignment(AssignmentDefinitionKind { assignment,
assignment: AstNodeRef::new(parsed.clone(), assignment), target_index,
target: AstNodeRef::new(parsed, target), name,
}) kind,
} }) => DefinitionKind::Assignment(AssignmentDefinitionKind {
assignment: AstNodeRef::new(parsed.clone(), assignment),
target_index,
name: AstNodeRef::new(parsed, name),
kind,
}),
DefinitionNodeRef::AnnotatedAssignment(assign) => { DefinitionNodeRef::AnnotatedAssignment(assign) => {
DefinitionKind::AnnotatedAssignment(AstNodeRef::new(parsed, assign)) DefinitionKind::AnnotatedAssignment(AstNodeRef::new(parsed, assign))
} }
@ -300,8 +307,10 @@ impl DefinitionNodeRef<'_> {
Self::NamedExpression(node) => node.into(), Self::NamedExpression(node) => node.into(),
Self::Assignment(AssignmentDefinitionNodeRef { Self::Assignment(AssignmentDefinitionNodeRef {
assignment: _, assignment: _,
target, target_index: _,
}) => target.into(), name,
kind: _,
}) => name.into(),
Self::AnnotatedAssignment(node) => node.into(), Self::AnnotatedAssignment(node) => node.into(),
Self::AugmentedAssignment(node) => node.into(), Self::AugmentedAssignment(node) => node.into(),
Self::For(ForStmtDefinitionNodeRef { Self::For(ForStmtDefinitionNodeRef {
@ -485,17 +494,34 @@ impl ImportFromDefinitionKind {
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct AssignmentDefinitionKind { pub struct AssignmentDefinitionKind {
assignment: AstNodeRef<ast::StmtAssign>, assignment: AstNodeRef<ast::StmtAssign>,
target: AstNodeRef<ast::ExprName>, target_index: usize,
name: AstNodeRef<ast::ExprName>,
kind: AssignmentKind,
} }
impl AssignmentDefinitionKind { impl AssignmentDefinitionKind {
pub(crate) fn assignment(&self) -> &ast::StmtAssign { pub(crate) fn value(&self) -> &ast::Expr {
self.assignment.node() &self.assignment.node().value
} }
pub(crate) fn target(&self) -> &ast::ExprName { pub(crate) fn target(&self) -> &ast::Expr {
self.target.node() &self.assignment.node().targets[self.target_index]
} }
pub(crate) fn name(&self) -> &ast::ExprName {
self.name.node()
}
pub(crate) fn kind(&self) -> AssignmentKind {
self.kind
}
}
/// The kind of assignment target expression.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AssignmentKind {
Sequence,
Name,
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]

View file

@ -1510,6 +1510,12 @@ pub struct StringLiteralType<'db> {
value: Box<str>, value: Box<str>,
} }
impl<'db> StringLiteralType<'db> {
pub fn len(&self, db: &'db dyn Db) -> usize {
self.value(db).len()
}
}
#[salsa::interned] #[salsa::interned]
pub struct BytesLiteralType<'db> { pub struct BytesLiteralType<'db> {
#[return_ref] #[return_ref]
@ -1522,6 +1528,16 @@ pub struct TupleType<'db> {
elements: Box<[Type<'db>]>, elements: Box<[Type<'db>]>,
} }
impl<'db> TupleType<'db> {
pub fn get(&self, db: &'db dyn Db, index: usize) -> Option<Type<'db>> {
self.elements(db).get(index).copied()
}
pub fn len(&self, db: &'db dyn Db) -> usize {
self.elements(db).len()
}
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::{ use super::{

View file

@ -27,6 +27,7 @@
//! associated with a particular definition. Scope-level inference infers deferred types for all //! associated with a particular definition. Scope-level inference infers deferred types for all
//! definitions once the rest of the types in the scope have been inferred. //! definitions once the rest of the types in the scope have been inferred.
use itertools::Itertools; use itertools::Itertools;
use std::borrow::Cow;
use std::num::NonZeroU32; use std::num::NonZeroU32;
use ruff_db::files::File; use ruff_db::files::File;
@ -41,7 +42,7 @@ use crate::module_name::ModuleName;
use crate::module_resolver::{file_to_module, resolve_module}; use crate::module_resolver::{file_to_module, resolve_module};
use crate::semantic_index::ast_ids::{HasScopedAstId, HasScopedUseId, ScopedExpressionId}; use crate::semantic_index::ast_ids::{HasScopedAstId, HasScopedUseId, ScopedExpressionId};
use crate::semantic_index::definition::{ use crate::semantic_index::definition::{
Definition, DefinitionKind, DefinitionNodeKey, ExceptHandlerDefinitionKind, AssignmentKind, Definition, DefinitionKind, DefinitionNodeKey, ExceptHandlerDefinitionKind,
}; };
use crate::semantic_index::expression::Expression; use crate::semantic_index::expression::Expression;
use crate::semantic_index::semantic_index; use crate::semantic_index::semantic_index;
@ -415,7 +416,9 @@ impl<'db> TypeInferenceBuilder<'db> {
DefinitionKind::Assignment(assignment) => { DefinitionKind::Assignment(assignment) => {
self.infer_assignment_definition( self.infer_assignment_definition(
assignment.target(), assignment.target(),
assignment.assignment(), assignment.value(),
assignment.name(),
assignment.kind(),
definition, definition,
); );
} }
@ -1151,13 +1154,23 @@ impl<'db> TypeInferenceBuilder<'db> {
} = assignment; } = assignment;
for target in targets { for target in targets {
if let ast::Expr::Name(name) = target { self.infer_assignment_target(target, value);
self.infer_definition(name); }
} else { }
// TODO infer definitions in unpacking assignment. When we do, this duplication of
// the "get `Expression`, call `infer_expression_types` on it, `self.extend`" dance // TODO: Remove the `value` argument once we handle all possible assignment targets.
// will be removed; it'll all happen in `infer_assignment_definition` instead. fn infer_assignment_target(&mut self, target: &ast::Expr, value: &ast::Expr) {
let expression = self.index.expression(value.as_ref()); match target {
ast::Expr::Name(name) => self.infer_definition(name),
ast::Expr::List(ast::ExprList { elts, .. })
| ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => {
for element in elts {
self.infer_assignment_target(element, value);
}
}
_ => {
// TODO: Remove this once we handle all possible assignment targets.
let expression = self.index.expression(value);
self.extend(infer_expression_types(self.db, expression)); self.extend(infer_expression_types(self.db, expression));
self.infer_expression(target); self.infer_expression(target);
} }
@ -1166,18 +1179,138 @@ impl<'db> TypeInferenceBuilder<'db> {
fn infer_assignment_definition( fn infer_assignment_definition(
&mut self, &mut self,
target: &ast::ExprName, target: &ast::Expr,
assignment: &ast::StmtAssign, value: &ast::Expr,
name: &ast::ExprName,
kind: AssignmentKind,
definition: Definition<'db>, definition: Definition<'db>,
) { ) {
let expression = self.index.expression(assignment.value.as_ref()); let expression = self.index.expression(value);
let result = infer_expression_types(self.db, expression); let result = infer_expression_types(self.db, expression);
self.extend(result); self.extend(result);
let value_ty = self.expression_ty(&assignment.value);
self.add_binding(assignment.into(), definition, value_ty); let value_ty = self.expression_ty(value);
let target_ty = match kind {
AssignmentKind::Sequence => self.infer_sequence_unpacking(target, value_ty, name),
AssignmentKind::Name => value_ty,
};
self.add_binding(name.into(), definition, target_ty);
self.types self.types
.expressions .expressions
.insert(target.scoped_ast_id(self.db, self.scope), value_ty); .insert(name.scoped_ast_id(self.db, self.scope), target_ty);
}
fn infer_sequence_unpacking(
&mut self,
target: &ast::Expr,
value_ty: Type<'db>,
name: &ast::ExprName,
) -> Type<'db> {
// The inner function is recursive and only differs in the return type which is an `Option`
// where if the variable is found, the corresponding type is returned otherwise `None`.
fn inner<'db>(
builder: &mut TypeInferenceBuilder<'db>,
target: &ast::Expr,
value_ty: Type<'db>,
name: &ast::ExprName,
) -> Option<Type<'db>> {
match target {
ast::Expr::Name(target_name) if target_name == name => {
return Some(value_ty);
}
ast::Expr::Starred(ast::ExprStarred { value, .. }) => {
return inner(builder, value, value_ty, name);
}
ast::Expr::List(ast::ExprList { elts, .. })
| ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => match value_ty {
Type::Tuple(tuple_ty) => {
let starred_index = elts.iter().position(ast::Expr::is_starred_expr);
let element_types = if let Some(starred_index) = starred_index {
if tuple_ty.len(builder.db) >= elts.len() - 1 {
let mut element_types = Vec::with_capacity(elts.len());
element_types.extend_from_slice(
// SAFETY: Safe because of the length check above.
&tuple_ty.elements(builder.db)[..starred_index],
);
// E.g., in `(a, *b, c, d) = ...`, the index of starred element `b`
// is 1 and the remaining elements after that are 2.
let remaining = elts.len() - (starred_index + 1);
// This index represents the type of the last element that belongs
// to the starred expression, in an exclusive manner.
let starred_end_index = tuple_ty.len(builder.db) - remaining;
// SAFETY: Safe because of the length check above.
let _starred_element_types = &tuple_ty.elements(builder.db)
[starred_index..starred_end_index];
// TODO: Combine the types into a list type. If the
// starred_element_types is empty, then it should be `List[Any]`.
// combine_types(starred_element_types);
element_types.push(Type::Todo);
element_types.extend_from_slice(
// SAFETY: Safe because of the length check above.
&tuple_ty.elements(builder.db)[starred_end_index..],
);
Cow::Owned(element_types)
} else {
let mut element_types = tuple_ty.elements(builder.db).to_vec();
element_types.insert(starred_index, Type::Todo);
Cow::Owned(element_types)
}
} else {
Cow::Borrowed(tuple_ty.elements(builder.db).as_ref())
};
for (index, element) in elts.iter().enumerate() {
if let Some(ty) = inner(
builder,
element,
element_types.get(index).copied().unwrap_or(Type::Unknown),
name,
) {
return Some(ty);
}
}
}
Type::StringLiteral(string_literal_ty) => {
// Deconstruct the string literal to delegate the inference back to the
// tuple type for correct handling of starred expressions. We could go
// further and deconstruct to an array of `StringLiteral` with each
// individual character, instead of just an array of `LiteralString`, but
// there would be a cost and it's not clear that it's worth it.
let value_ty = Type::Tuple(TupleType::new(
builder.db,
vec![Type::LiteralString; string_literal_ty.len(builder.db)]
.into_boxed_slice(),
));
if let Some(ty) = inner(builder, target, value_ty, name) {
return Some(ty);
}
}
_ => {
let value_ty = if matches!(value_ty, Type::LiteralString) {
Type::LiteralString
} else {
value_ty
.iterate(builder.db)
.unwrap_with_diagnostic(AnyNodeRef::from(target), builder)
};
for element in elts {
if let Some(ty) = inner(builder, element, value_ty, name) {
return Some(ty);
}
}
}
},
_ => {}
}
None
}
inner(self, target, value_ty, name).unwrap_or(Type::Unknown)
} }
fn infer_annotated_assignment_statement(&mut self, assignment: &ast::StmtAnnAssign) { fn infer_annotated_assignment_statement(&mut self, assignment: &ast::StmtAnnAssign) {

View file

@ -9,7 +9,7 @@ use ruff_db::parsed::parsed_module;
use ruff_db::system::{OsSystem, SystemPath, SystemPathBuf}; use ruff_db::system::{OsSystem, SystemPath, SystemPathBuf};
use ruff_python_ast::visitor::source_order; use ruff_python_ast::visitor::source_order;
use ruff_python_ast::visitor::source_order::SourceOrderVisitor; use ruff_python_ast::visitor::source_order::SourceOrderVisitor;
use ruff_python_ast::{Alias, Expr, Parameter, ParameterWithDefault, Stmt}; use ruff_python_ast::{self as ast, Alias, Expr, Parameter, ParameterWithDefault, Stmt};
fn setup_db(workspace_root: &SystemPath) -> anyhow::Result<RootDatabase> { fn setup_db(workspace_root: &SystemPath) -> anyhow::Result<RootDatabase> {
let system = OsSystem::new(workspace_root); let system = OsSystem::new(workspace_root);
@ -65,6 +65,17 @@ impl<'db> PullTypesVisitor<'db> {
model: SemanticModel::new(db, file), model: SemanticModel::new(db, file),
} }
} }
fn visit_assign_target(&mut self, target: &Expr) {
match target {
Expr::List(ast::ExprList { elts, .. }) | Expr::Tuple(ast::ExprTuple { elts, .. }) => {
for element in elts {
self.visit_assign_target(element);
}
}
_ => self.visit_expr(target),
}
}
} }
impl SourceOrderVisitor<'_> for PullTypesVisitor<'_> { impl SourceOrderVisitor<'_> for PullTypesVisitor<'_> {
@ -76,10 +87,15 @@ impl SourceOrderVisitor<'_> for PullTypesVisitor<'_> {
Stmt::ClassDef(class) => { Stmt::ClassDef(class) => {
let _ty = class.ty(&self.model); let _ty = class.ty(&self.model);
} }
Stmt::Assign(assign) => {
for target in &assign.targets {
self.visit_assign_target(target);
}
return;
}
Stmt::AnnAssign(_) Stmt::AnnAssign(_)
| Stmt::Return(_) | Stmt::Return(_)
| Stmt::Delete(_) | Stmt::Delete(_)
| Stmt::Assign(_)
| Stmt::AugAssign(_) | Stmt::AugAssign(_)
| Stmt::TypeAlias(_) | Stmt::TypeAlias(_)
| Stmt::For(_) | Stmt::For(_)