Skip to content

Commit

Permalink
Don't skip over imports and other nodes containing nested statements …
Browse files Browse the repository at this point in the history
…in import collector (#13521)
  • Loading branch information
MichaReiser authored Sep 26, 2024
1 parent 9442cd8 commit ff2d214
Show file tree
Hide file tree
Showing 3 changed files with 97 additions and 33 deletions.
55 changes: 55 additions & 0 deletions crates/ruff/tests/analyze_graph.rs
Original file line number Diff line number Diff line change
Expand Up @@ -367,3 +367,58 @@ fn wildcard() -> Result<()> {

Ok(())
}

#[test]
fn nested_imports() -> Result<()> {
let tempdir = TempDir::new()?;
let root = ChildPath::new(tempdir.path());

root.child("ruff").child("__init__.py").write_str("")?;
root.child("ruff")
.child("a.py")
.write_str(indoc::indoc! {r#"
match x:
case 1:
import ruff.b
"#})?;
root.child("ruff")
.child("b.py")
.write_str(indoc::indoc! {r#"
try:
import ruff.c
except ImportError as e:
import ruff.d
"#})?;
root.child("ruff")
.child("c.py")
.write_str(indoc::indoc! {r#"def c(): ..."#})?;
root.child("ruff")
.child("d.py")
.write_str(indoc::indoc! {r#"def d(): ..."#})?;

insta::with_settings!({
filters => INSTA_FILTERS.to_vec(),
}, {
assert_cmd_snapshot!(command().current_dir(&root), @r#"
success: true
exit_code: 0
----- stdout -----
{
"ruff/__init__.py": [],
"ruff/a.py": [
"ruff/b.py"
],
"ruff/b.py": [
"ruff/c.py",
"ruff/d.py"
],
"ruff/c.py": [],
"ruff/d.py": []
}
----- stderr -----
"#);
});

Ok(())
}
57 changes: 32 additions & 25 deletions crates/ruff_graph/src/collector.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
use red_knot_python_semantic::ModuleName;
use ruff_python_ast::visitor::source_order::{
walk_expr, walk_module, walk_stmt, SourceOrderVisitor, TraversalSignal,
walk_expr, walk_module, walk_stmt, SourceOrderVisitor,
};
use ruff_python_ast::{self as ast, AnyNodeRef, Expr, Mod, Stmt};
use ruff_python_ast::{self as ast, Expr, Mod, Stmt};

/// Collect all imports for a given Python file.
#[derive(Default, Debug)]
Expand Down Expand Up @@ -32,28 +32,6 @@ impl<'a> Collector<'a> {
}

impl<'ast> SourceOrderVisitor<'ast> for Collector<'_> {
fn enter_node(&mut self, node: AnyNodeRef<'ast>) -> TraversalSignal {
// If string detection is enabled, we have to visit everything. Otherwise, we should only
// visit compounds statements, which can contain import statements.
if self.string_imports
|| matches!(
node,
AnyNodeRef::ModModule(_)
| AnyNodeRef::StmtFunctionDef(_)
| AnyNodeRef::StmtClassDef(_)
| AnyNodeRef::StmtWhile(_)
| AnyNodeRef::StmtFor(_)
| AnyNodeRef::StmtWith(_)
| AnyNodeRef::StmtIf(_)
| AnyNodeRef::StmtTry(_)
)
{
TraversalSignal::Traverse
} else {
TraversalSignal::Skip
}
}

fn visit_stmt(&mut self, stmt: &'ast Stmt) {
match stmt {
Stmt::ImportFrom(ast::StmtImportFrom {
Expand Down Expand Up @@ -107,9 +85,38 @@ impl<'ast> SourceOrderVisitor<'ast> for Collector<'_> {
}
}
}
_ => {
Stmt::FunctionDef(_)
| Stmt::ClassDef(_)
| Stmt::While(_)
| Stmt::If(_)
| Stmt::With(_)
| Stmt::Match(_)
| Stmt::Try(_)
| Stmt::For(_) => {
// Always traverse into compound statements.
walk_stmt(self, stmt);
}

Stmt::Return(_)
| Stmt::Delete(_)
| Stmt::Assign(_)
| Stmt::AugAssign(_)
| Stmt::AnnAssign(_)
| Stmt::TypeAlias(_)
| Stmt::Raise(_)
| Stmt::Assert(_)
| Stmt::Global(_)
| Stmt::Nonlocal(_)
| Stmt::Expr(_)
| Stmt::Pass(_)
| Stmt::Break(_)
| Stmt::Continue(_)
| Stmt::IpyEscapeCommand(_) => {
// Only traverse simple statements when string imports is enabled.
if self.string_imports {
walk_stmt(self, stmt);
}
}
}
}

Expand Down
18 changes: 10 additions & 8 deletions crates/ruff_graph/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,21 +1,23 @@
use crate::collector::Collector;
pub use crate::db::ModuleDb;
use crate::resolver::Resolver;
pub use crate::settings::{AnalyzeSettings, Direction};
use std::collections::{BTreeMap, BTreeSet};

use anyhow::Result;

use ruff_db::system::{SystemPath, SystemPathBuf};
use ruff_python_ast::helpers::to_module_path;
use ruff_python_parser::{parse, Mode};
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, BTreeSet};

use crate::collector::Collector;
pub use crate::db::ModuleDb;
use crate::resolver::Resolver;
pub use crate::settings::{AnalyzeSettings, Direction};

mod collector;
mod db;
mod resolver;
mod settings;

#[derive(Debug, Default)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ModuleImports(BTreeSet<SystemPathBuf>);

impl ModuleImports {
Expand Down Expand Up @@ -90,7 +92,7 @@ impl ModuleImports {
}

#[derive(Debug, Default)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ImportMap(BTreeMap<SystemPathBuf, ModuleImports>);

impl ImportMap {
Expand Down

0 comments on commit ff2d214

Please sign in to comment.