Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support relative imports in banned-api enforcement #4025

Merged
merged 1 commit into from
Apr 19, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 16 additions & 6 deletions crates/ruff/src/checkers/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
);
}
}
Expand Down
11 changes: 5 additions & 6 deletions crates/ruff/src/rules/flake8_tidy_imports/banned_api.rs
Original file line number Diff line number Diff line change
@@ -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};

Expand Down Expand Up @@ -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<T>(checker: &mut Checker, name: String, located: &Located<T>) {
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),
));
}
}
Expand Down
8 changes: 8 additions & 0 deletions crates/ruff/src/rules/flake8_tidy_imports/relative_imports.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)) {
Expand All @@ -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)) {
Expand Down
119 changes: 116 additions & 3 deletions crates/ruff_python_ast/src/helpers.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use std::borrow::Cow;
use std::path::Path;

use itertools::Itertools;
Expand Down Expand Up @@ -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<usize>, module: Option<&str>) -> String {
let mut module_name = String::with_capacity(16);
if let Some(level) = level {
Expand All @@ -677,6 +688,16 @@ pub fn format_import_from(level: Option<usize>, 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<usize>,
module: Option<&str>,
Expand Down Expand Up @@ -715,7 +736,24 @@ pub fn to_module_path(package: &Path, path: &Path) -> Option<Vec<String>> {
.collect::<Option<Vec<String>>>()
}

/// 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);

Expand All @@ -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();
}

Expand All @@ -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<usize>,
module: Option<&'a str>,
module_path: Option<&[String]>,
) -> Option<Cow<'a, str>> {
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> {
Expand Down Expand Up @@ -1350,13 +1424,15 @@ pub fn locate_cmpops(contents: &str) -> Vec<LocatedCmpop> {

#[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;
Expand Down Expand Up @@ -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#"
Expand Down
16 changes: 12 additions & 4 deletions crates/ruff_python_semantic/src/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand All @@ -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
}
Expand Down