[typing] Add find_assigned_value helper func to typing.rs to retrieve value of a given variable id (#8583)

qdegraaf · web-flow · commit 8314c8bb0565 · 2023-12-13T00:24:47.000Z
## Summary

Adds `find_assigned_value` a function which gets the `&amp;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
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_trio/TRIO115.py b/crates/ruff_linter/resources/test/fixtures/flake8_trio/TRIO115.py
@@ -19,6 +19,28 @@ async def func():
     bar = "bar"
     trio.sleep(bar)
 
+    x, y = 0, 2000
+    trio.sleep(x)  # TRIO115
+    trio.sleep(y)  # OK
+
+    (a, b, [c, (d, e)]) = (1, 2, (0, [4, 0]))
+    trio.sleep(c)  # TRIO115
+    trio.sleep(d)  # OK
+    trio.sleep(e)  # TRIO115
+
+    m_x, m_y = 0
+    trio.sleep(m_y)  # TRIO115
+    trio.sleep(m_x)  # TRIO115
+
+    m_a = m_b = 0
+    trio.sleep(m_a)  # TRIO115
+    trio.sleep(m_b)  # TRIO115
+
+    m_c = (m_d, m_e) = (0, 0)
+    trio.sleep(m_c)  # OK
+    trio.sleep(m_d)  # TRIO115
+    trio.sleep(m_e)  # TRIO115
+
 
 def func():
     trio.run(trio.sleep(0))  # TRIO115
diff --git a/crates/ruff_linter/resources/test/fixtures/perflint/PERF101.py b/crates/ruff_linter/resources/test/fixtures/perflint/PERF101.py
@@ -63,3 +63,8 @@
 
 for i in list(foo_set):  # Ok
     foo_set.append(i + 1)
+
+x, y, nested_tuple = (1, 2, (3, 4, 5))
+
+for i in list(nested_tuple):  # PERF101
+    pass
diff --git a/crates/ruff_linter/src/rules/flake8_trio/rules/zero_sleep_call.rs b/crates/ruff_linter/src/rules/flake8_trio/rules/zero_sleep_call.rs
@@ -1,7 +1,7 @@
 use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::Stmt;
 use ruff_python_ast::{self as ast, Expr, ExprCall, Int};
+use ruff_python_semantic::analyze::typing::find_assigned_value;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -71,30 +71,15 @@ pub(crate) fn zero_sleep_call(checker: &mut Checker, call: &ExprCall) {
             }
         }
         Expr::Name(ast::ExprName { id, .. }) => {
-            let scope = checker.semantic().current_scope();
-            if let Some(binding_id) = scope.get(id) {
-                let binding = checker.semantic().binding(binding_id);
-                if binding.kind.is_assignment() || binding.kind.is_named_expr_assignment() {
-                    if let Some(parent_id) = binding.source {
-                        let parent = checker.semantic().statement(parent_id);
-                        if let Stmt::Assign(ast::StmtAssign { value, .. })
-                        | Stmt::AnnAssign(ast::StmtAnnAssign {
-                            value: Some(value), ..
-                        })
-                        | Stmt::AugAssign(ast::StmtAugAssign { value, .. }) = parent
-                        {
-                            let Expr::NumberLiteral(ast::ExprNumberLiteral { value: num, .. }) =
-                                value.as_ref()
-                            else {
-                                return;
-                            };
-                            let Some(int) = num.as_int() else { return };
-                            if *int != Int::ZERO {
-                                return;
-                            }
-                        }
-                    }
-                }
+            let Some(value) = find_assigned_value(id, checker.semantic()) else {
+                return;
+            };
+            let Expr::NumberLiteral(ast::ExprNumberLiteral { value: num, .. }) = value else {
+                return;
+            };
+            let Some(int) = num.as_int() else { return };
+            if *int != Int::ZERO {
+                return;
             }
         }
         _ => return,
diff --git a/crates/ruff_linter/src/rules/flake8_trio/snapshots/ruff_linter__rules__flake8_trio__tests__TRIO115_TRIO115.py.snap b/crates/ruff_linter/src/rules/flake8_trio/snapshots/ruff_linter__rules__flake8_trio__tests__TRIO115_TRIO115.py.snap
@@ -85,51 +85,230 @@ TRIO115.py:17:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.s
 19 19 |     bar = "bar"
 20 20 |     trio.sleep(bar)
 
-TRIO115.py:31:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+TRIO115.py:23:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
    |
-30 | def func():
-31 |     sleep(0)  # TRIO115
-   |     ^^^^^^^^ TRIO115
+22 |     x, y = 0, 2000
+23 |     trio.sleep(x)  # TRIO115
+   |     ^^^^^^^^^^^^^ TRIO115
+24 |     trio.sleep(y)  # OK
    |
    = help: Replace with `trio.lowlevel.checkpoint()`
 
 ℹ Safe fix
-24 24 |     trio.run(trio.sleep(0))  # TRIO115
+20 20 |     trio.sleep(bar)
+21 21 | 
+22 22 |     x, y = 0, 2000
+23    |-    trio.sleep(x)  # TRIO115
+   23 |+    trio.lowlevel.checkpoint()  # TRIO115
+24 24 |     trio.sleep(y)  # OK
 25 25 | 
-26 26 | 
-27    |-from trio import Event, sleep
-   27 |+from trio import Event, sleep, lowlevel
-28 28 | 
-29 29 | 
-30 30 | def func():
-31    |-    sleep(0)  # TRIO115
-   31 |+    lowlevel.checkpoint()  # TRIO115
-32 32 | 
-33 33 | 
-34 34 | async def func():
-
-TRIO115.py:35:11: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
-   |
-34 | async def func():
-35 |     await sleep(seconds=0)  # TRIO115
-   |           ^^^^^^^^^^^^^^^^ TRIO115
+26 26 |     (a, b, [c, (d, e)]) = (1, 2, (0, [4, 0]))
+
+TRIO115.py:27:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+   |
+26 |     (a, b, [c, (d, e)]) = (1, 2, (0, [4, 0]))
+27 |     trio.sleep(c)  # TRIO115
+   |     ^^^^^^^^^^^^^ TRIO115
+28 |     trio.sleep(d)  # OK
+29 |     trio.sleep(e)  # TRIO115
    |
    = help: Replace with `trio.lowlevel.checkpoint()`
 
 ℹ Safe fix
-24 24 |     trio.run(trio.sleep(0))  # TRIO115
+24 24 |     trio.sleep(y)  # OK
 25 25 | 
-26 26 | 
-27    |-from trio import Event, sleep
-   27 |+from trio import Event, sleep, lowlevel
-28 28 | 
-29 29 | 
-30 30 | def func():
+26 26 |     (a, b, [c, (d, e)]) = (1, 2, (0, [4, 0]))
+27    |-    trio.sleep(c)  # TRIO115
+   27 |+    trio.lowlevel.checkpoint()  # TRIO115
+28 28 |     trio.sleep(d)  # OK
+29 29 |     trio.sleep(e)  # TRIO115
+30 30 | 
+
+TRIO115.py:29:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+   |
+27 |     trio.sleep(c)  # TRIO115
+28 |     trio.sleep(d)  # OK
+29 |     trio.sleep(e)  # TRIO115
+   |     ^^^^^^^^^^^^^ TRIO115
+30 | 
+31 |     m_x, m_y = 0
+   |
+   = help: Replace with `trio.lowlevel.checkpoint()`
+
+ℹ Safe fix
+26 26 |     (a, b, [c, (d, e)]) = (1, 2, (0, [4, 0]))
+27 27 |     trio.sleep(c)  # TRIO115
+28 28 |     trio.sleep(d)  # OK
+29    |-    trio.sleep(e)  # TRIO115
+   29 |+    trio.lowlevel.checkpoint()  # TRIO115
+30 30 | 
+31 31 |     m_x, m_y = 0
+32 32 |     trio.sleep(m_y)  # TRIO115
+
+TRIO115.py:32:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+   |
+31 |     m_x, m_y = 0
+32 |     trio.sleep(m_y)  # TRIO115
+   |     ^^^^^^^^^^^^^^^ TRIO115
+33 |     trio.sleep(m_x)  # TRIO115
+   |
+   = help: Replace with `trio.lowlevel.checkpoint()`
+
+ℹ Safe fix
+29 29 |     trio.sleep(e)  # TRIO115
+30 30 | 
+31 31 |     m_x, m_y = 0
+32    |-    trio.sleep(m_y)  # TRIO115
+   32 |+    trio.lowlevel.checkpoint()  # TRIO115
+33 33 |     trio.sleep(m_x)  # TRIO115
+34 34 | 
+35 35 |     m_a = m_b = 0
+
+TRIO115.py:33:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+   |
+31 |     m_x, m_y = 0
+32 |     trio.sleep(m_y)  # TRIO115
+33 |     trio.sleep(m_x)  # TRIO115
+   |     ^^^^^^^^^^^^^^^ TRIO115
+34 | 
+35 |     m_a = m_b = 0
+   |
+   = help: Replace with `trio.lowlevel.checkpoint()`
+
+ℹ Safe fix
+30 30 | 
+31 31 |     m_x, m_y = 0
+32 32 |     trio.sleep(m_y)  # TRIO115
+33    |-    trio.sleep(m_x)  # TRIO115
+   33 |+    trio.lowlevel.checkpoint()  # TRIO115
+34 34 | 
+35 35 |     m_a = m_b = 0
+36 36 |     trio.sleep(m_a)  # TRIO115
+
+TRIO115.py:36:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+   |
+35 |     m_a = m_b = 0
+36 |     trio.sleep(m_a)  # TRIO115
+   |     ^^^^^^^^^^^^^^^ TRIO115
+37 |     trio.sleep(m_b)  # TRIO115
+   |
+   = help: Replace with `trio.lowlevel.checkpoint()`
+
+ℹ Safe fix
+33 33 |     trio.sleep(m_x)  # TRIO115
+34 34 | 
+35 35 |     m_a = m_b = 0
+36    |-    trio.sleep(m_a)  # TRIO115
+   36 |+    trio.lowlevel.checkpoint()  # TRIO115
+37 37 |     trio.sleep(m_b)  # TRIO115
+38 38 | 
+39 39 |     m_c = (m_d, m_e) = (0, 0)
+
+TRIO115.py:37:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+   |
+35 |     m_a = m_b = 0
+36 |     trio.sleep(m_a)  # TRIO115
+37 |     trio.sleep(m_b)  # TRIO115
+   |     ^^^^^^^^^^^^^^^ TRIO115
+38 | 
+39 |     m_c = (m_d, m_e) = (0, 0)
+   |
+   = help: Replace with `trio.lowlevel.checkpoint()`
+
+ℹ Safe fix
+34 34 | 
+35 35 |     m_a = m_b = 0
+36 36 |     trio.sleep(m_a)  # TRIO115
+37    |-    trio.sleep(m_b)  # TRIO115
+   37 |+    trio.lowlevel.checkpoint()  # TRIO115
+38 38 | 
+39 39 |     m_c = (m_d, m_e) = (0, 0)
+40 40 |     trio.sleep(m_c)  # OK
+
+TRIO115.py:41:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+   |
+39 |     m_c = (m_d, m_e) = (0, 0)
+40 |     trio.sleep(m_c)  # OK
+41 |     trio.sleep(m_d)  # TRIO115
+   |     ^^^^^^^^^^^^^^^ TRIO115
+42 |     trio.sleep(m_e)  # TRIO115
+   |
+   = help: Replace with `trio.lowlevel.checkpoint()`
+
+ℹ Safe fix
+38 38 | 
+39 39 |     m_c = (m_d, m_e) = (0, 0)
+40 40 |     trio.sleep(m_c)  # OK
+41    |-    trio.sleep(m_d)  # TRIO115
+   41 |+    trio.lowlevel.checkpoint()  # TRIO115
+42 42 |     trio.sleep(m_e)  # TRIO115
+43 43 | 
+44 44 | 
+
+TRIO115.py:42:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+   |
+40 |     trio.sleep(m_c)  # OK
+41 |     trio.sleep(m_d)  # TRIO115
+42 |     trio.sleep(m_e)  # TRIO115
+   |     ^^^^^^^^^^^^^^^ TRIO115
+   |
+   = help: Replace with `trio.lowlevel.checkpoint()`
+
+ℹ Safe fix
+39 39 |     m_c = (m_d, m_e) = (0, 0)
+40 40 |     trio.sleep(m_c)  # OK
+41 41 |     trio.sleep(m_d)  # TRIO115
+42    |-    trio.sleep(m_e)  # TRIO115
+   42 |+    trio.lowlevel.checkpoint()  # TRIO115
+43 43 | 
+44 44 | 
+45 45 | def func():
+
+TRIO115.py:53:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+   |
+52 | def func():
+53 |     sleep(0)  # TRIO115
+   |     ^^^^^^^^ TRIO115
+   |
+   = help: Replace with `trio.lowlevel.checkpoint()`
+
+ℹ Safe fix
+46 46 |     trio.run(trio.sleep(0))  # TRIO115
+47 47 | 
+48 48 | 
+49    |-from trio import Event, sleep
+   49 |+from trio import Event, sleep, lowlevel
+50 50 | 
+51 51 | 
+52 52 | def func():
+53    |-    sleep(0)  # TRIO115
+   53 |+    lowlevel.checkpoint()  # TRIO115
+54 54 | 
+55 55 | 
+56 56 | async def func():
+
+TRIO115.py:57:11: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+   |
+56 | async def func():
+57 |     await sleep(seconds=0)  # TRIO115
+   |           ^^^^^^^^^^^^^^^^ TRIO115
+   |
+   = help: Replace with `trio.lowlevel.checkpoint()`
+
+ℹ Safe fix
+46 46 |     trio.run(trio.sleep(0))  # TRIO115
+47 47 | 
+48 48 | 
+49    |-from trio import Event, sleep
+   49 |+from trio import Event, sleep, lowlevel
+50 50 | 
+51 51 | 
+52 52 | def func():
 --------------------------------------------------------------------------------
-32 32 | 
-33 33 | 
-34 34 | async def func():
-35    |-    await sleep(seconds=0)  # TRIO115
-   35 |+    await lowlevel.checkpoint()  # TRIO115
+54 54 | 
+55 55 | 
+56 56 | async def func():
+57    |-    await sleep(seconds=0)  # TRIO115
+   57 |+    await lowlevel.checkpoint()  # TRIO115
 
 
diff --git a/crates/ruff_linter/src/rules/perflint/rules/unnecessary_list_cast.rs b/crates/ruff_linter/src/rules/perflint/rules/unnecessary_list_cast.rs
@@ -1,7 +1,7 @@
 use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::Stmt;
-use ruff_python_ast::{self as ast, Arguments, Expr};
+use ruff_python_ast::{self as ast, Arguments, Expr, Stmt};
+use ruff_python_semantic::analyze::typing::find_assigned_value;
 use ruff_text_size::TextRange;
 
 use crate::checkers::ast::Checker;
@@ -110,30 +110,13 @@ pub(crate) fn unnecessary_list_cast(checker: &mut Checker, iter: &Expr, body: &[
             if body.iter().any(|stmt| match_append(stmt, id)) {
                 return;
             }
-            let scope = checker.semantic().current_scope();
-            if let Some(binding_id) = scope.get(id) {
-                let binding = checker.semantic().binding(binding_id);
-                if binding.kind.is_assignment() || binding.kind.is_named_expr_assignment() {
-                    if let Some(parent_id) = binding.source {
-                        let parent = checker.semantic().statement(parent_id);
-                        if let Stmt::Assign(ast::StmtAssign { value, .. })
-                        | Stmt::AnnAssign(ast::StmtAnnAssign {
-                            value: Some(value), ..
-                        })
-                        | Stmt::AugAssign(ast::StmtAugAssign { value, .. }) = parent
-                        {
-                            if matches!(
-                                value.as_ref(),
-                                Expr::Tuple(_) | Expr::List(_) | Expr::Set(_)
-                            ) {
-                                let mut diagnostic =
-                                    Diagnostic::new(UnnecessaryListCast, *list_range);
-                                diagnostic.set_fix(remove_cast(*list_range, *iterable_range));
-                                checker.diagnostics.push(diagnostic);
-                            }
-                        }
-                    }
-                }
+            let Some(value) = find_assigned_value(id, checker.semantic()) else {
+                return;
+            };
+            if matches!(value, Expr::Tuple(_) | Expr::List(_) | Expr::Set(_)) {
+                let mut diagnostic = Diagnostic::new(UnnecessaryListCast, *list_range);
+                diagnostic.set_fix(remove_cast(*list_range, *iterable_range));
+                checker.diagnostics.push(diagnostic);
             }
         }
         _ => {}
diff --git a/crates/ruff_linter/src/rules/perflint/snapshots/ruff_linter__rules__perflint__tests__PERF101_PERF101.py.snap b/crates/ruff_linter/src/rules/perflint/snapshots/ruff_linter__rules__perflint__tests__PERF101_PERF101.py.snap
diff --git a/crates/ruff_python_semantic/src/analyze/typing.rs b/crates/ruff_python_semantic/src/analyze/typing.rs
