Skip to content

Commit

Permalink
Avoid adding required imports to stub files (#3940)
Browse files Browse the repository at this point in the history
  • Loading branch information
charliermarsh authored Apr 12, 2023
1 parent 61200d2 commit eb0dd74
Show file tree
Hide file tree
Showing 10 changed files with 83 additions and 10 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
"""Hello, world!"""

x = 1
7 changes: 5 additions & 2 deletions crates/ruff/src/checkers/imports.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ use ruff_python_ast::helpers::to_module_path;
use ruff_python_ast::imports::{ImportMap, ModuleImport};
use ruff_python_ast::source_code::{Indexer, Locator, Stylist};
use ruff_python_ast::visitor::Visitor;
use ruff_python_stdlib::path::is_python_stub_file;

use crate::directives::IsortDirectives;
use crate::registry::Rule;
Expand Down Expand Up @@ -88,9 +89,11 @@ pub fn check_imports(
path: &Path,
package: Option<&Path>,
) -> (Vec<Diagnostic>, Option<ImportMap>) {
let is_stub = is_python_stub_file(path);

// Extract all imports from the AST.
let tracker = {
let mut tracker = ImportTracker::new(locator, directives, path);
let mut tracker = ImportTracker::new(locator, directives, is_stub);
tracker.visit_body(python_ast);
tracker
};
Expand All @@ -111,7 +114,7 @@ pub fn check_imports(
}
if settings.rules.enabled(Rule::MissingRequiredImport) {
diagnostics.extend(isort::rules::add_required_imports(
&blocks, python_ast, locator, stylist, settings, autofix,
&blocks, python_ast, locator, stylist, settings, autofix, is_stub,
));
}

Expand Down
4 changes: 4 additions & 0 deletions crates/ruff/src/rules/isort/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -714,6 +714,7 @@ mod tests {
}

#[test_case(Path::new("docstring.py"))]
#[test_case(Path::new("docstring.pyi"))]
#[test_case(Path::new("docstring_only.py"))]
#[test_case(Path::new("multiline_docstring.py"))]
#[test_case(Path::new("empty.py"))]
Expand All @@ -737,6 +738,7 @@ mod tests {
}

#[test_case(Path::new("docstring.py"))]
#[test_case(Path::new("docstring.pyi"))]
#[test_case(Path::new("docstring_only.py"))]
#[test_case(Path::new("empty.py"))]
fn required_imports(path: &Path) -> Result<()> {
Expand All @@ -760,6 +762,7 @@ mod tests {
}

#[test_case(Path::new("docstring.py"))]
#[test_case(Path::new("docstring.pyi"))]
#[test_case(Path::new("docstring_only.py"))]
#[test_case(Path::new("empty.py"))]
fn combined_required_imports(path: &Path) -> Result<()> {
Expand All @@ -782,6 +785,7 @@ mod tests {
}

#[test_case(Path::new("docstring.py"))]
#[test_case(Path::new("docstring.pyi"))]
#[test_case(Path::new("docstring_only.py"))]
#[test_case(Path::new("empty.py"))]
fn straight_required_import(path: &Path) -> Result<()> {
Expand Down
12 changes: 11 additions & 1 deletion crates/ruff/src/rules/isort/rules/add_required_imports.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use rustpython_parser::ast::{Location, StmtKind, Suite};
use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::helpers::is_docstring_stmt;
use ruff_python_ast::imports::{Alias, AnyImport, Import, ImportFrom};
use ruff_python_ast::imports::{Alias, AnyImport, FutureImport, Import, ImportFrom};
use ruff_python_ast::source_code::{Locator, Stylist};
use ruff_python_ast::types::Range;

Expand Down Expand Up @@ -84,6 +84,7 @@ fn contains(block: &Block, required_import: &AnyImport) -> bool {
})
}

#[allow(clippy::too_many_arguments)]
fn add_required_import(
required_import: &AnyImport,
blocks: &[&Block],
Expand All @@ -92,6 +93,7 @@ fn add_required_import(
stylist: &Stylist,
settings: &Settings,
autofix: flags::Autofix,
is_stub: bool,
) -> Option<Diagnostic> {
// If the import is already present in a top-level block, don't add it.
if blocks
Expand All @@ -107,6 +109,11 @@ fn add_required_import(
return None;
}

// We don't need to add `__future__` imports to stubs.
if is_stub && required_import.is_future_import() {
return None;
}

// Always insert the diagnostic at top-of-file.
let mut diagnostic = Diagnostic::new(
MissingRequiredImport(required_import.to_string()),
Expand All @@ -126,6 +133,7 @@ pub fn add_required_imports(
stylist: &Stylist,
settings: &Settings,
autofix: flags::Autofix,
is_stub: bool,
) -> Vec<Diagnostic> {
settings
.isort
Expand Down Expand Up @@ -167,6 +175,7 @@ pub fn add_required_imports(
stylist,
settings,
autofix,
is_stub,
)
})
.collect(),
Expand All @@ -186,6 +195,7 @@ pub fn add_required_imports(
stylist,
settings,
autofix,
is_stub,
)
})
.collect(),
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
---
source: crates/ruff/src/rules/isort/mod.rs
---

Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
---
source: crates/ruff/src/rules/isort/mod.rs
---

Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
---
source: crates/ruff/src/rules/isort/mod.rs
---

Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
---
source: crates/ruff/src/rules/isort/mod.rs
---
docstring.pyi:1:1: I002 [*] Missing required import: `import os`
|
1 | """Hello, world!"""
| I002
2 |
3 | x = 1
|
= help: Insert required import: `import os`

Suggested fix
1 1 | """Hello, world!"""
2 |+import os
2 3 |
3 4 | x = 1


10 changes: 3 additions & 7 deletions crates/ruff/src/rules/isort/track.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
use std::path::Path;

use rustpython_parser::ast::{
Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Constant, Excepthandler,
ExcepthandlerKind, Expr, ExprContext, Keyword, MatchCase, Operator, Pattern, Stmt, StmtKind,
Expand All @@ -8,11 +6,9 @@ use rustpython_parser::ast::{

use ruff_python_ast::source_code::Locator;
use ruff_python_ast::visitor::Visitor;
use ruff_python_stdlib::path::is_python_stub_file;

use crate::directives::IsortDirectives;

use super::helpers;
use crate::rules::isort::helpers;

#[derive(Debug, Copy, Clone)]
pub enum Trailer {
Expand All @@ -38,11 +34,11 @@ pub struct ImportTracker<'a> {
}

impl<'a> ImportTracker<'a> {
pub fn new(locator: &'a Locator<'a>, directives: &'a IsortDirectives, path: &'a Path) -> Self {
pub fn new(locator: &'a Locator<'a>, directives: &'a IsortDirectives, is_stub: bool) -> Self {
Self {
locator,
directives,
is_stub: is_python_stub_file(path),
is_stub,
blocks: vec![Block::default()],
split_index: 0,
nested: false,
Expand Down
26 changes: 26 additions & 0 deletions crates/ruff_python_ast/src/imports.rs
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,32 @@ impl std::fmt::Display for ImportFrom<'_> {
}
}

pub trait FutureImport {
/// Returns `true` if this import is from the `__future__` module.
fn is_future_import(&self) -> bool;
}

impl FutureImport for Import<'_> {
fn is_future_import(&self) -> bool {
self.name.name == "__future__"
}
}

impl FutureImport for ImportFrom<'_> {
fn is_future_import(&self) -> bool {
self.module == Some("__future__")
}
}

impl FutureImport for AnyImport<'_> {
fn is_future_import(&self) -> bool {
match self {
AnyImport::Import(import) => import.is_future_import(),
AnyImport::ImportFrom(import_from) => import_from.is_future_import(),
}
}
}

/// A representation of a module reference in an import statement.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ModuleImport {
Expand Down

0 comments on commit eb0dd74

Please sign in to comment.