Skip to content

Commit

Permalink
Update corpus tests to visit assignment targets
Browse files Browse the repository at this point in the history
  • Loading branch information
dhruvmanila committed Oct 14, 2024
1 parent 81f31b0 commit c8dbb21
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ reveal_type(b) # revealed: Unknown
### Simple unpacking

```py
a, b = 'abc'
a, b = 'ab'
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
```
Expand Down
20 changes: 18 additions & 2 deletions crates/red_knot_workspace/tests/check.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use ruff_db::parsed::parsed_module;
use ruff_db::system::{OsSystem, SystemPath, SystemPathBuf};
use ruff_python_ast::visitor::source_order;
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> {
let system = OsSystem::new(workspace_root);
Expand Down Expand Up @@ -65,6 +65,17 @@ impl<'db> PullTypesVisitor<'db> {
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<'_> {
Expand All @@ -76,10 +87,15 @@ impl SourceOrderVisitor<'_> for PullTypesVisitor<'_> {
Stmt::ClassDef(class) => {
let _ty = class.ty(&self.model);
}
Stmt::Assign(assign) => {
for target in &assign.targets {
self.visit_assign_target(target);
}
return;
}
Stmt::AnnAssign(_)
| Stmt::Return(_)
| Stmt::Delete(_)
| Stmt::Assign(_)
| Stmt::AugAssign(_)
| Stmt::TypeAlias(_)
| Stmt::For(_)
Expand Down

0 comments on commit c8dbb21

Please sign in to comment.