[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:
qdegraaf 2023-12-13 01:24:47 +01:00 committed by GitHub
parent cb201bc4a5
commit 8314c8bb05
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 380 additions and 83 deletions

View file

@ -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
}