mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-30 13:51:37 +00:00
[typing
] Add find_assigned_value
helper func to typing.rs
to retrieve value of a given variable id
(#8583)
## Summary Adds `find_assigned_value` a function which gets the `&Expr` assigned to a given `id` if one exists in the semantic model. Open TODOs: - [ ] Handle `binding.kind.is_unpacked_assignment()`: I am bit confused by this one. The snippet from its documentation does not appear to be counted as an unpacked assignment and the only ones I could find for which that was true were invalid Python like: ```python x, y = 1 ``` - [ ] How to handle AugAssign. Can we combine statements like: ```python (a, b) = [(1, 2, 3), (4,)] a += (6, 7) ``` to get the full value for a? Code currently just returns `None` for these assign types - [ ] Multi target assigns ```python m_c = (m_d, m_e) = (0, 0) trio.sleep(m_c) # OK trio.sleep(m_d) # TRIO115 trio.sleep(m_e) # TRIO115 ``` ## Test Plan Used the function in two rules: - `TRIO115` - `PERF101` Expanded both their fixtures for explicit multi target check
This commit is contained in:
parent
cb201bc4a5
commit
8314c8bb05
7 changed files with 380 additions and 83 deletions
|
@ -568,3 +568,108 @@ pub fn resolve_assignment<'a>(
|
|||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Find the assigned [`Expr`] for a given symbol, if any.
|
||||
///
|
||||
/// For example given:
|
||||
/// ```python
|
||||
/// foo = 42
|
||||
/// (bar, bla) = 1, "str"
|
||||
/// ```
|
||||
///
|
||||
/// This function will return a `NumberLiteral` with value `Int(42)` when called with `foo` and a
|
||||
/// `StringLiteral` with value `"str"` when called with `bla`.
|
||||
pub fn find_assigned_value<'a>(symbol: &str, semantic: &'a SemanticModel<'a>) -> Option<&'a Expr> {
|
||||
let binding_id = semantic.lookup_symbol(symbol)?;
|
||||
let binding = semantic.binding(binding_id);
|
||||
if binding.kind.is_assignment() || binding.kind.is_named_expr_assignment() {
|
||||
let parent_id = binding.source?;
|
||||
let parent = semantic.statement(parent_id);
|
||||
match parent {
|
||||
Stmt::Assign(ast::StmtAssign { value, targets, .. }) => match value.as_ref() {
|
||||
Expr::Tuple(ast::ExprTuple { elts, .. })
|
||||
| Expr::List(ast::ExprList { elts, .. }) => {
|
||||
if let Some(target) = targets.iter().find(|target| defines(symbol, target)) {
|
||||
return match target {
|
||||
Expr::Tuple(ast::ExprTuple {
|
||||
elts: target_elts, ..
|
||||
})
|
||||
| Expr::List(ast::ExprList {
|
||||
elts: target_elts, ..
|
||||
})
|
||||
| Expr::Set(ast::ExprSet {
|
||||
elts: target_elts, ..
|
||||
}) => get_value_by_id(symbol, target_elts, elts),
|
||||
_ => Some(value.as_ref()),
|
||||
};
|
||||
}
|
||||
}
|
||||
_ => return Some(value.as_ref()),
|
||||
},
|
||||
Stmt::AnnAssign(ast::StmtAnnAssign {
|
||||
value: Some(value), ..
|
||||
}) => {
|
||||
return Some(value.as_ref());
|
||||
}
|
||||
Stmt::AugAssign(_) => return None,
|
||||
_ => return None,
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Returns `true` if the [`Expr`] defines the symbol.
|
||||
fn defines(symbol: &str, expr: &Expr) -> bool {
|
||||
match expr {
|
||||
Expr::Name(ast::ExprName { id, .. }) => id == symbol,
|
||||
Expr::Tuple(ast::ExprTuple { elts, .. })
|
||||
| Expr::List(ast::ExprList { elts, .. })
|
||||
| Expr::Set(ast::ExprSet { elts, .. }) => elts.iter().any(|elt| defines(symbol, elt)),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn get_value_by_id<'a>(
|
||||
target_id: &str,
|
||||
targets: &'a [Expr],
|
||||
values: &'a [Expr],
|
||||
) -> Option<&'a Expr> {
|
||||
for (target, value) in targets.iter().zip(values.iter()) {
|
||||
match target {
|
||||
Expr::Tuple(ast::ExprTuple {
|
||||
elts: target_elts, ..
|
||||
})
|
||||
| Expr::List(ast::ExprList {
|
||||
elts: target_elts, ..
|
||||
})
|
||||
| Expr::Set(ast::ExprSet {
|
||||
elts: target_elts, ..
|
||||
}) => {
|
||||
// Collection types can be mismatched like in: (a, b, [c, d]) = [1, 2, {3, 4}]
|
||||
match value {
|
||||
Expr::Tuple(ast::ExprTuple {
|
||||
elts: value_elts, ..
|
||||
})
|
||||
| Expr::List(ast::ExprList {
|
||||
elts: value_elts, ..
|
||||
})
|
||||
| Expr::Set(ast::ExprSet {
|
||||
elts: value_elts, ..
|
||||
}) => {
|
||||
if let Some(result) = get_value_by_id(target_id, target_elts, value_elts) {
|
||||
return Some(result);
|
||||
}
|
||||
}
|
||||
_ => (),
|
||||
};
|
||||
}
|
||||
Expr::Name(ast::ExprName { id, .. }) => {
|
||||
if *id == target_id {
|
||||
return Some(value);
|
||||
}
|
||||
}
|
||||
_ => (),
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue