From 7fa1da20fba31410ee26ae1cf7233a9278c41638 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 19 Apr 2023 14:30:13 -0400 Subject: [PATCH] Support relative imports in `banned-api` enforcement (#4025) --- crates/ruff/src/checkers/ast/mod.rs | 22 +++- .../rules/flake8_tidy_imports/banned_api.rs | 11 +- .../flake8_tidy_imports/relative_imports.rs | 8 ++ crates/ruff_python_ast/src/helpers.rs | 119 +++++++++++++++++- crates/ruff_python_semantic/src/context.rs | 16 ++- 5 files changed, 157 insertions(+), 19 deletions(-) diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs index fc6692286b5a1..6c22449dcebca 100644 --- a/crates/ruff/src/checkers/ast/mod.rs +++ b/crates/ruff/src/checkers/ast/mod.rs @@ -1155,13 +1155,23 @@ where } if self.settings.rules.enabled(Rule::BannedApi) { - if level.map_or(true, |level| level == 0) { - if let Some(module) = module { - for name in names { - flake8_tidy_imports::banned_api::name_is_banned(self, module, name); + if let Some(module) = helpers::resolve_imported_module_path( + *level, + module.as_deref(), + self.module_path.as_deref(), + ) { + flake8_tidy_imports::banned_api::name_or_parent_is_banned( + self, &module, stmt, + ); + + for alias in names { + if alias.node.name == "*" { + continue; } - flake8_tidy_imports::banned_api::name_or_parent_is_banned( - self, module, stmt, + flake8_tidy_imports::banned_api::name_is_banned( + self, + format!("{module}.{}", alias.node.name), + alias, ); } } diff --git a/crates/ruff/src/rules/flake8_tidy_imports/banned_api.rs b/crates/ruff/src/rules/flake8_tidy_imports/banned_api.rs index 0671954b6874b..66dbf0c523f34 100644 --- a/crates/ruff/src/rules/flake8_tidy_imports/banned_api.rs +++ b/crates/ruff/src/rules/flake8_tidy_imports/banned_api.rs @@ -1,5 +1,5 @@ use rustc_hash::FxHashMap; -use rustpython_parser::ast::{Alias, Expr, Located}; +use rustpython_parser::ast::{Expr, Located}; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; @@ -51,16 +51,15 @@ impl Violation for BannedApi { } /// TID251 -pub fn name_is_banned(checker: &mut Checker, module: &str, name: &Alias) { +pub fn name_is_banned(checker: &mut Checker, name: String, located: &Located) { let banned_api = &checker.settings.flake8_tidy_imports.banned_api; - let full_name = format!("{module}.{}", &name.node.name); - if let Some(ban) = banned_api.get(&full_name) { + if let Some(ban) = banned_api.get(&name) { checker.diagnostics.push(Diagnostic::new( BannedApi { - name: full_name, + name, message: ban.msg.to_string(), }, - Range::from(name), + Range::from(located), )); } } diff --git a/crates/ruff/src/rules/flake8_tidy_imports/relative_imports.rs b/crates/ruff/src/rules/flake8_tidy_imports/relative_imports.rs index a660516f494c5..e70eec0080d42 100644 --- a/crates/ruff/src/rules/flake8_tidy_imports/relative_imports.rs +++ b/crates/ruff/src/rules/flake8_tidy_imports/relative_imports.rs @@ -106,6 +106,10 @@ fn fix_banned_relative_import( let module_name = if let Some(module) = module { let call_path = from_relative_import(&parts, module); + // Empty indicates an invalid module. + if call_path.is_empty() { + return None; + } // Require import to be a valid module: // https://python.org/dev/peps/pep-0008/#package-and-module-names if !call_path.iter().all(|part| is_identifier(part)) { @@ -115,6 +119,10 @@ fn fix_banned_relative_import( } else if parts.len() > 1 { let module = parts.pop().unwrap(); let call_path = from_relative_import(&parts, &module); + // Empty indicates an invalid module. + if call_path.is_empty() { + return None; + } // Require import to be a valid module: // https://python.org/dev/peps/pep-0008/#package-and-module-names if !call_path.iter().all(|part| is_identifier(part)) { diff --git a/crates/ruff_python_ast/src/helpers.rs b/crates/ruff_python_ast/src/helpers.rs index 8eed8fc45a125..7ee6e96fdcafc 100644 --- a/crates/ruff_python_ast/src/helpers.rs +++ b/crates/ruff_python_ast/src/helpers.rs @@ -1,3 +1,4 @@ +use std::borrow::Cow; use std::path::Path; use itertools::Itertools; @@ -662,7 +663,17 @@ where }) } -/// Format the module name for a relative import. +/// Format the module reference name for a relative import. +/// +/// # Examples +/// +/// ```rust +/// # use ruff_python_ast::helpers::format_import_from; +/// +/// assert_eq!(format_import_from(None, None), "".to_string()); +/// assert_eq!(format_import_from(Some(1), None), ".".to_string()); +/// assert_eq!(format_import_from(Some(1), Some("foo")), ".foo".to_string()); +/// ``` pub fn format_import_from(level: Option, module: Option<&str>) -> String { let mut module_name = String::with_capacity(16); if let Some(level) = level { @@ -677,6 +688,16 @@ pub fn format_import_from(level: Option, module: Option<&str>) -> String } /// Format the member reference name for a relative import. +/// +/// # Examples +/// +/// ```rust +/// # use ruff_python_ast::helpers::format_import_from_member; +/// +/// assert_eq!(format_import_from_member(None, None, "bar"), "bar".to_string()); +/// assert_eq!(format_import_from_member(Some(1), None, "bar"), ".bar".to_string()); +/// assert_eq!(format_import_from_member(Some(1), Some("foo"), "bar"), ".foo.bar".to_string()); +/// ``` pub fn format_import_from_member( level: Option, module: Option<&str>, @@ -715,7 +736,24 @@ pub fn to_module_path(package: &Path, path: &Path) -> Option> { .collect::>>() } -/// Create a call path from a relative import. +/// Create a [`CallPath`] from a relative import reference name (like `".foo.bar"`). +/// +/// Returns an empty [`CallPath`] if the import is invalid (e.g., a relative import that +/// extends beyond the top-level module). +/// +/// # Examples +/// +/// ```rust +/// # use smallvec::{smallvec, SmallVec}; +/// # use ruff_python_ast::helpers::from_relative_import; +/// +/// assert_eq!(from_relative_import(&[], "bar"), SmallVec::from_buf(["bar"])); +/// assert_eq!(from_relative_import(&["foo".to_string()], "bar"), SmallVec::from_buf(["foo", "bar"])); +/// assert_eq!(from_relative_import(&["foo".to_string()], "bar.baz"), SmallVec::from_buf(["foo", "bar", "baz"])); +/// assert_eq!(from_relative_import(&["foo".to_string()], ".bar"), SmallVec::from_buf(["bar"])); +/// assert!(from_relative_import(&["foo".to_string()], "..bar").is_empty()); +/// assert!(from_relative_import(&["foo".to_string()], "...bar").is_empty()); +/// ``` pub fn from_relative_import<'a>(module: &'a [String], name: &'a str) -> CallPath<'a> { let mut call_path: CallPath = SmallVec::with_capacity(module.len() + 1); @@ -724,6 +762,9 @@ pub fn from_relative_import<'a>(module: &'a [String], name: &'a str) -> CallPath // Remove segments based on the number of dots. for _ in 0..name.chars().take_while(|c| *c == '.').count() { + if call_path.is_empty() { + return SmallVec::new(); + } call_path.pop(); } @@ -733,6 +774,39 @@ pub fn from_relative_import<'a>(module: &'a [String], name: &'a str) -> CallPath call_path } +/// Given an imported module (based on its relative import level and module name), return the +/// fully-qualified module path. +pub fn resolve_imported_module_path<'a>( + level: Option, + module: Option<&'a str>, + module_path: Option<&[String]>, +) -> Option> { + let Some(level) = level else { + return Some(Cow::Borrowed(module.unwrap_or(""))); + }; + + if level == 0 { + return Some(Cow::Borrowed(module.unwrap_or(""))); + } + + let Some(module_path) = module_path else { + return None; + }; + + if level >= module_path.len() { + return None; + } + + let mut qualified_path = module_path[..module_path.len() - level].join("."); + if let Some(module) = module { + if !qualified_path.is_empty() { + qualified_path.push('.'); + } + qualified_path.push_str(module); + } + Some(Cow::Owned(qualified_path)) +} + /// A [`Visitor`] that collects all `return` statements in a function or method. #[derive(Default)] pub struct ReturnStatementVisitor<'a> { @@ -1350,13 +1424,15 @@ pub fn locate_cmpops(contents: &str) -> Vec { #[cfg(test)] mod tests { + use std::borrow::Cow; + use anyhow::Result; use rustpython_parser as parser; use rustpython_parser::ast::{Cmpop, Location}; use crate::helpers::{ elif_else_range, else_range, first_colon_range, identifier_range, locate_cmpops, - match_trailing_content, LocatedCmpop, + match_trailing_content, resolve_imported_module_path, LocatedCmpop, }; use crate::source_code::Locator; use crate::types::Range; @@ -1469,6 +1545,43 @@ class Class(): Ok(()) } + #[test] + fn resolve_import() { + // Return the module directly. + assert_eq!( + resolve_imported_module_path(None, Some("foo"), None), + Some(Cow::Borrowed("foo")) + ); + + // Construct the module path from the calling module's path. + assert_eq!( + resolve_imported_module_path( + Some(1), + Some("foo"), + Some(&["bar".to_string(), "baz".to_string()]) + ), + Some(Cow::Owned("bar.foo".to_string())) + ); + + // We can't return the module if it's a relative import, and we don't know the calling + // module's path. + assert_eq!( + resolve_imported_module_path(Some(1), Some("foo"), None), + None + ); + + // We can't return the module if it's a relative import, and the path goes beyond the + // calling module's path. + assert_eq!( + resolve_imported_module_path(Some(1), Some("foo"), Some(&["bar".to_string()])), + None, + ); + assert_eq!( + resolve_imported_module_path(Some(2), Some("foo"), Some(&["bar".to_string()])), + None + ); + } + #[test] fn extract_else_range() -> Result<()> { let contents = r#" diff --git a/crates/ruff_python_semantic/src/context.rs b/crates/ruff_python_semantic/src/context.rs index ed6f73a932e0c..1d67601b52b71 100644 --- a/crates/ruff_python_semantic/src/context.rs +++ b/crates/ruff_python_semantic/src/context.rs @@ -174,8 +174,12 @@ impl<'a> Context<'a> { if name.starts_with('.') { if let Some(module) = &self.module_path { let mut source_path = from_relative_import(module, name); - source_path.extend(call_path.into_iter().skip(1)); - Some(source_path) + if source_path.is_empty() { + None + } else { + source_path.extend(call_path.into_iter().skip(1)); + Some(source_path) + } } else { None } @@ -191,8 +195,12 @@ impl<'a> Context<'a> { if name.starts_with('.') { if let Some(module) = &self.module_path { let mut source_path = from_relative_import(module, name); - source_path.extend(call_path.into_iter().skip(1)); - Some(source_path) + if source_path.is_empty() { + None + } else { + source_path.extend(call_path.into_iter().skip(1)); + Some(source_path) + } } else { None }