From 76148ddb76396602c57da141c12f83396c9b754e Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sat, 5 Aug 2023 11:21:50 -0400 Subject: [PATCH] Store call paths rather than stringified names (#6102) ## Summary Historically, we've stored "qualified names" on our `BindingKind::Import`, `BindingKind::SubmoduleImport`, and `BindingKind::ImportFrom` structs. In Ruff, a "qualified name" is a dot-separated path to a symbol. For example, given `import foo.bar`, the "qualified name" would be `"foo.bar"`; and given `from foo.bar import baz`, the "qualified name" would be `foo.bar.baz`. This PR modifies the `BindingKind` structs to instead store _call paths_ rather than qualified names. So in the examples above, we'd store `["foo", "bar"]` and `["foo", "bar", "baz"]`. It turns out that this more efficient given our data access patterns. Namely, we frequently need to convert the qualified name to a call path (whenever we call `resolve_call_path`), and it turns out that we do this operation enough that those conversations show up on benchmarks. There are a few other advantages to using call paths, rather than qualified names: 1. The size of `BindingKind` is reduced from 32 to 24 bytes, since we no longer need to store a `String` (only a boxed slice). 2. All three import types are more consistent, since they now all store a boxed slice, rather than some storing an `&str` and some storing a `String` (for `BindingKind::ImportFrom`, we needed to allocate a `String` to create the qualified name, but the call path is a slice of static elements that don't require that allocation). 3. A lot of code gets simpler, in part because we now do call path resolution "earlier". Most notably, for relative imports (`from .foo import bar`), we store the _resolved_ call path rather than the relative call path, so the semantic model doesn't have to deal with that resolution. (See that `resolve_call_path` is simpler, fewer branches, etc.) In my testing, this change improves the all-rules benchmark by another 4-5% on top of the improvements mentioned in #6047. --- crates/ruff/src/autofix/codemods.rs | 85 ++------- crates/ruff/src/autofix/edits.rs | 4 +- crates/ruff/src/checkers/ast/mod.rs | 22 ++- crates/ruff/src/importer/mod.rs | 20 +- crates/ruff/src/renamer.rs | 2 +- .../rules/unconventional_import_alias.rs | 10 +- .../unaliased_collections_abc_set_import.rs | 7 +- .../runtime_import_in_type_checking_block.rs | 47 ++--- .../rules/typing_only_runtime_import.rs | 104 +++++------ crates/ruff/src/rules/pandas_vet/helpers.rs | 41 ++--- .../pandas_vet/rules/inplace_argument.rs | 18 +- .../src/rules/pyflakes/rules/unused_import.rs | 43 +++-- .../src/rules/pyupgrade/rules/f_strings.rs | 5 +- .../rules/pyupgrade/rules/os_error_alias.rs | 2 +- .../rules/unnecessary_builtin_import.rs | 9 +- .../rules/unnecessary_future_import.rs | 9 +- crates/ruff_python_ast/src/call_path.rs | 32 +++- crates/ruff_python_ast/src/helpers.rs | 79 ++++++-- crates/ruff_python_semantic/src/binding.rs | 173 +++++++++++++----- crates/ruff_python_semantic/src/model.rs | 102 +++++------ 20 files changed, 449 insertions(+), 365 deletions(-) diff --git a/crates/ruff/src/autofix/codemods.rs b/crates/ruff/src/autofix/codemods.rs index 18f7723b3c..be7028d5f3 100644 --- a/crates/ruff/src/autofix/codemods.rs +++ b/crates/ruff/src/autofix/codemods.rs @@ -33,7 +33,7 @@ impl<'a, T: Codegen<'a>> CodegenStylist<'a> for T { /// /// Returns `Ok(None)` if the statement is empty after removing the imports. pub(crate) fn remove_imports<'a>( - imports: impl Iterator, + member_names: impl Iterator, stmt: &Stmt, locator: &Locator, stylist: &Stylist, @@ -45,27 +45,20 @@ pub(crate) fn remove_imports<'a>( bail!("Expected Statement::Simple"); }; - let (aliases, import_module) = match body.body.first_mut() { - Some(SmallStatement::Import(import_body)) => (&mut import_body.names, None), + let aliases = match body.body.first_mut() { + Some(SmallStatement::Import(import_body)) => &mut import_body.names, Some(SmallStatement::ImportFrom(import_body)) => { if let ImportNames::Aliases(names) = &mut import_body.names { - ( - names, - Some((&import_body.relative, import_body.module.as_ref())), - ) + names } else if let ImportNames::Star(..) = &import_body.names { // Special-case: if the import is a `from ... import *`, then we delete the // entire statement. let mut found_star = false; - for import in imports { - let qualified_name = match import_body.module.as_ref() { - Some(module_name) => format!("{}.*", compose_module_path(module_name)), - None => "*".to_string(), - }; - if import == qualified_name { + for member in member_names { + if member == "*" { found_star = true; } else { - bail!("Expected \"*\" for unused import (got: \"{}\")", import); + bail!("Expected \"*\" for unused import (got: \"{}\")", member); } } if !found_star { @@ -82,30 +75,10 @@ pub(crate) fn remove_imports<'a>( // Preserve the trailing comma (or not) from the last entry. let trailing_comma = aliases.last().and_then(|alias| alias.comma.clone()); - for import in imports { - let alias_index = aliases.iter().position(|alias| { - let qualified_name = match import_module { - Some((relative, module)) => { - let module = module.map(compose_module_path); - let member = compose_module_path(&alias.name); - let mut qualified_name = String::with_capacity( - relative.len() + module.as_ref().map_or(0, String::len) + member.len() + 1, - ); - for _ in 0..relative.len() { - qualified_name.push('.'); - } - if let Some(module) = module { - qualified_name.push_str(&module); - qualified_name.push('.'); - } - qualified_name.push_str(&member); - qualified_name - } - None => compose_module_path(&alias.name), - }; - qualified_name == import - }); - + for member in member_names { + let alias_index = aliases + .iter() + .position(|alias| member == compose_module_path(&alias.name)); if let Some(index) = alias_index { aliases.remove(index); } @@ -139,7 +112,7 @@ pub(crate) fn remove_imports<'a>( /// /// Returns the modified import statement. pub(crate) fn retain_imports( - imports: &[&str], + member_names: &[&str], stmt: &Stmt, locator: &Locator, stylist: &Stylist, @@ -151,14 +124,11 @@ pub(crate) fn retain_imports( bail!("Expected Statement::Simple"); }; - let (aliases, import_module) = match body.body.first_mut() { - Some(SmallStatement::Import(import_body)) => (&mut import_body.names, None), + let aliases = match body.body.first_mut() { + Some(SmallStatement::Import(import_body)) => &mut import_body.names, Some(SmallStatement::ImportFrom(import_body)) => { if let ImportNames::Aliases(names) = &mut import_body.names { - ( - names, - Some((&import_body.relative, import_body.module.as_ref())), - ) + names } else { bail!("Expected: ImportNames::Aliases"); } @@ -170,28 +140,9 @@ pub(crate) fn retain_imports( let trailing_comma = aliases.last().and_then(|alias| alias.comma.clone()); aliases.retain(|alias| { - imports.iter().any(|import| { - let qualified_name = match import_module { - Some((relative, module)) => { - let module = module.map(compose_module_path); - let member = compose_module_path(&alias.name); - let mut qualified_name = String::with_capacity( - relative.len() + module.as_ref().map_or(0, String::len) + member.len() + 1, - ); - for _ in 0..relative.len() { - qualified_name.push('.'); - } - if let Some(module) = module { - qualified_name.push_str(&module); - qualified_name.push('.'); - } - qualified_name.push_str(&member); - qualified_name - } - None => compose_module_path(&alias.name), - }; - qualified_name == *import - }) + member_names + .iter() + .any(|member| *member == compose_module_path(&alias.name)) }); // But avoid destroying any trailing comments. diff --git a/crates/ruff/src/autofix/edits.rs b/crates/ruff/src/autofix/edits.rs index 6e90771d8e..17736349a5 100644 --- a/crates/ruff/src/autofix/edits.rs +++ b/crates/ruff/src/autofix/edits.rs @@ -58,14 +58,14 @@ pub(crate) fn delete_stmt( /// Generate a `Fix` to remove the specified imports from an `import` statement. pub(crate) fn remove_unused_imports<'a>( - unused_imports: impl Iterator, + member_names: impl Iterator, stmt: &Stmt, parent: Option<&Stmt>, locator: &Locator, stylist: &Stylist, indexer: &Indexer, ) -> Result { - match codemods::remove_imports(unused_imports, stmt, locator, stylist)? { + match codemods::remove_imports(member_names, stmt, locator, stylist)? { None => Ok(delete_stmt(stmt, parent, locator, indexer)), Some(content) => Ok(Edit::range_replacement(content, stmt.range())), } diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs index 1c19eda25e..6f02558e14 100644 --- a/crates/ruff/src/checkers/ast/mod.rs +++ b/crates/ruff/src/checkers/ast/mod.rs @@ -39,7 +39,9 @@ use ruff_text_size::{TextRange, TextSize}; use ruff_diagnostics::{Diagnostic, IsolationLevel}; use ruff_python_ast::all::{extract_all_names, DunderAllFlags}; -use ruff_python_ast::helpers::{extract_handled_exceptions, to_module_path}; +use ruff_python_ast::helpers::{ + collect_import_from_member, extract_handled_exceptions, to_module_path, +}; use ruff_python_ast::identifier::Identifier; use ruff_python_ast::str::trailing_quote; use ruff_python_ast::visitor::{walk_except_handler, walk_pattern, Visitor}; @@ -320,11 +322,11 @@ where // Given `import foo.bar`, `name` would be "foo", and `qualified_name` would be // "foo.bar". let name = alias.name.split('.').next().unwrap(); - let qualified_name = &alias.name; + let call_path: Box<[&str]> = alias.name.split('.').collect(); self.add_binding( name, alias.identifier(), - BindingKind::SubmoduleImport(SubmoduleImport { qualified_name }), + BindingKind::SubmoduleImport(SubmoduleImport { call_path }), BindingFlags::EXTERNAL, ); } else { @@ -341,11 +343,11 @@ where } let name = alias.asname.as_ref().unwrap_or(&alias.name); - let qualified_name = &alias.name; + let call_path: Box<[&str]> = alias.name.split('.').collect(); self.add_binding( name, alias.identifier(), - BindingKind::Import(Import { qualified_name }), + BindingKind::Import(Import { call_path }), flags, ); } @@ -389,12 +391,16 @@ where // be "foo.bar". Given `from foo import bar as baz`, `name` would be "baz" // and `qualified_name` would be "foo.bar". let name = alias.asname.as_ref().unwrap_or(&alias.name); - let qualified_name = - helpers::format_import_from_member(level, module, &alias.name); + + // Attempt to resolve any relative imports; but if we don't know the current + // module path, or the relative import extends beyond the package root, + // fallback to a literal representation (e.g., `[".", "foo"]`). + let call_path = collect_import_from_member(level, module, &alias.name) + .into_boxed_slice(); self.add_binding( name, alias.identifier(), - BindingKind::FromImport(FromImport { qualified_name }), + BindingKind::FromImport(FromImport { call_path }), flags, ); } diff --git a/crates/ruff/src/importer/mod.rs b/crates/ruff/src/importer/mod.rs index 1ef341ee7e..0fcdff9b55 100644 --- a/crates/ruff/src/importer/mod.rs +++ b/crates/ruff/src/importer/mod.rs @@ -87,13 +87,13 @@ impl<'a> Importer<'a> { /// import statement. pub(crate) fn runtime_import_edit( &self, - import: &StmtImports, + import: &ImportedMembers, at: TextSize, ) -> Result { // Generate the modified import statement. let content = autofix::codemods::retain_imports( - &import.qualified_names, - import.stmt, + &import.names, + import.statement, self.locator, self.stylist, )?; @@ -118,15 +118,15 @@ impl<'a> Importer<'a> { /// `TYPE_CHECKING` block. pub(crate) fn typing_import_edit( &self, - import: &StmtImports, + import: &ImportedMembers, at: TextSize, semantic: &SemanticModel, source_type: PySourceType, ) -> Result { // Generate the modified import statement. let content = autofix::codemods::retain_imports( - &import.qualified_names, - import.stmt, + &import.names, + import.statement, self.locator, self.stylist, )?; @@ -452,11 +452,11 @@ impl<'a> ImportRequest<'a> { } /// An existing list of module or member imports, located within an import statement. -pub(crate) struct StmtImports<'a> { +pub(crate) struct ImportedMembers<'a> { /// The import statement. - pub(crate) stmt: &'a Stmt, - /// The "qualified names" of the imported modules or members. - pub(crate) qualified_names: Vec<&'a str>, + pub(crate) statement: &'a Stmt, + /// The "names" of the imported members. + pub(crate) names: Vec<&'a str>, } /// The result of an [`Importer::get_or_import_symbol`] call. diff --git a/crates/ruff/src/renamer.rs b/crates/ruff/src/renamer.rs index 84f4cfbe76..d3f3e6fc92 100644 --- a/crates/ruff/src/renamer.rs +++ b/crates/ruff/src/renamer.rs @@ -231,7 +231,7 @@ impl Renamer { } BindingKind::SubmoduleImport(import) => { // Ex) Rename `import pandas.core` to `import pandas as pd`. - let module_name = import.qualified_name.split('.').next().unwrap(); + let module_name = import.call_path.first().unwrap(); Some(Edit::range_replacement( format!("{module_name} as {target}"), binding.range, diff --git a/crates/ruff/src/rules/flake8_import_conventions/rules/unconventional_import_alias.rs b/crates/ruff/src/rules/flake8_import_conventions/rules/unconventional_import_alias.rs index 628e7b167c..6a009f3834 100644 --- a/crates/ruff/src/rules/flake8_import_conventions/rules/unconventional_import_alias.rs +++ b/crates/ruff/src/rules/flake8_import_conventions/rules/unconventional_import_alias.rs @@ -2,7 +2,7 @@ use rustc_hash::FxHashMap; use ruff_diagnostics::{AutofixKind, Diagnostic, Fix, Violation}; use ruff_macros::{derive_message_formats, violation}; -use ruff_python_semantic::Binding; +use ruff_python_semantic::{Binding, Imported}; use crate::checkers::ast::Checker; use crate::registry::AsRule; @@ -56,11 +56,13 @@ pub(crate) fn unconventional_import_alias( binding: &Binding, conventions: &FxHashMap, ) -> Option { - let Some(qualified_name) = binding.qualified_name() else { + let Some(import) = binding.as_any_import() else { return None; }; - let Some(expected_alias) = conventions.get(qualified_name) else { + let qualified_name = import.qualified_name(); + + let Some(expected_alias) = conventions.get(qualified_name.as_str()) else { return None; }; @@ -71,7 +73,7 @@ pub(crate) fn unconventional_import_alias( let mut diagnostic = Diagnostic::new( UnconventionalImportAlias { - name: qualified_name.to_string(), + name: qualified_name, asname: expected_alias.to_string(), }, binding.range, diff --git a/crates/ruff/src/rules/flake8_pyi/rules/unaliased_collections_abc_set_import.rs b/crates/ruff/src/rules/flake8_pyi/rules/unaliased_collections_abc_set_import.rs index 6afd64b25b..90f6481255 100644 --- a/crates/ruff/src/rules/flake8_pyi/rules/unaliased_collections_abc_set_import.rs +++ b/crates/ruff/src/rules/flake8_pyi/rules/unaliased_collections_abc_set_import.rs @@ -1,6 +1,7 @@ use ruff_diagnostics::{AutofixKind, Diagnostic, Fix, Violation}; use ruff_macros::{derive_message_formats, violation}; -use ruff_python_semantic::{Binding, BindingKind, FromImport}; +use ruff_python_semantic::Imported; +use ruff_python_semantic::{Binding, BindingKind}; use crate::checkers::ast::Checker; use crate::registry::AsRule; @@ -50,10 +51,10 @@ pub(crate) fn unaliased_collections_abc_set_import( checker: &Checker, binding: &Binding, ) -> Option { - let BindingKind::FromImport(FromImport { qualified_name }) = &binding.kind else { + let BindingKind::FromImport(import) = &binding.kind else { return None; }; - if qualified_name.as_str() != "collections.abc.Set" { + if !matches!(import.call_path(), ["collections", "abc", "Set"]) { return None; } diff --git a/crates/ruff/src/rules/flake8_type_checking/rules/runtime_import_in_type_checking_block.rs b/crates/ruff/src/rules/flake8_type_checking/rules/runtime_import_in_type_checking_block.rs index 9bf591a5b5..66c187fc55 100644 --- a/crates/ruff/src/rules/flake8_type_checking/rules/runtime_import_in_type_checking_block.rs +++ b/crates/ruff/src/rules/flake8_type_checking/rules/runtime_import_in_type_checking_block.rs @@ -1,15 +1,16 @@ use anyhow::Result; use ruff_text_size::TextRange; use rustc_hash::FxHashMap; +use std::borrow::Cow; use ruff_diagnostics::{AutofixKind, Diagnostic, Fix, Violation}; use ruff_macros::{derive_message_formats, violation}; -use ruff_python_semantic::{NodeId, ResolvedReferenceId, Scope}; +use ruff_python_semantic::{AnyImport, Imported, NodeId, ResolvedReferenceId, Scope}; use crate::autofix; use crate::checkers::ast::Checker; use crate::codes::Rule; -use crate::importer::StmtImports; +use crate::importer::ImportedMembers; /// ## What it does /// Checks for runtime imports defined in a type-checking block. @@ -69,13 +70,13 @@ pub(crate) fn runtime_import_in_type_checking_block( diagnostics: &mut Vec, ) { // Collect all runtime imports by statement. - let mut errors_by_statement: FxHashMap> = FxHashMap::default(); - let mut ignores_by_statement: FxHashMap> = FxHashMap::default(); + let mut errors_by_statement: FxHashMap> = FxHashMap::default(); + let mut ignores_by_statement: FxHashMap> = FxHashMap::default(); for binding_id in scope.binding_ids() { let binding = checker.semantic().binding(binding_id); - let Some(qualified_name) = binding.qualified_name() else { + let Some(import) = binding.as_any_import() else { continue; }; @@ -96,8 +97,8 @@ pub(crate) fn runtime_import_in_type_checking_block( continue; }; - let import = Import { - qualified_name, + let import = ImportBinding { + import, reference_id, range: binding.range, parent_range: binding.parent_range(checker.semantic()), @@ -130,8 +131,8 @@ pub(crate) fn runtime_import_in_type_checking_block( None }; - for Import { - qualified_name, + for ImportBinding { + import, range, parent_range, .. @@ -139,7 +140,7 @@ pub(crate) fn runtime_import_in_type_checking_block( { let mut diagnostic = Diagnostic::new( RuntimeImportInTypeCheckingBlock { - qualified_name: qualified_name.to_string(), + qualified_name: import.qualified_name(), }, range, ); @@ -155,8 +156,8 @@ pub(crate) fn runtime_import_in_type_checking_block( // Separately, generate a diagnostic for every _ignored_ import, to ensure that the // suppression comments aren't marked as unused. - for Import { - qualified_name, + for ImportBinding { + import, range, parent_range, .. @@ -164,7 +165,7 @@ pub(crate) fn runtime_import_in_type_checking_block( { let mut diagnostic = Diagnostic::new( RuntimeImportInTypeCheckingBlock { - qualified_name: qualified_name.to_string(), + qualified_name: import.qualified_name(), }, range, ); @@ -176,9 +177,9 @@ pub(crate) fn runtime_import_in_type_checking_block( } /// A runtime-required import with its surrounding context. -struct Import<'a> { +struct ImportBinding<'a> { /// The qualified name of the import (e.g., `typing.List` for `from typing import List`). - qualified_name: &'a str, + import: AnyImport<'a>, /// The first reference to the imported symbol. reference_id: ResolvedReferenceId, /// The trimmed range of the import (e.g., `List` in `from typing import List`). @@ -188,18 +189,18 @@ struct Import<'a> { } /// Generate a [`Fix`] to remove runtime imports from a type-checking block. -fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[Import]) -> Result { +fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[ImportBinding]) -> Result { let stmt = checker.semantic().stmts[stmt_id]; let parent = checker.semantic().stmts.parent(stmt); - let qualified_names: Vec<&str> = imports + let member_names: Vec> = imports .iter() - .map(|Import { qualified_name, .. }| *qualified_name) + .map(|ImportBinding { import, .. }| import.member_name()) .collect(); // Find the first reference across all imports. let at = imports .iter() - .map(|Import { reference_id, .. }| { + .map(|ImportBinding { reference_id, .. }| { checker.semantic().reference(*reference_id).range().start() }) .min() @@ -207,7 +208,7 @@ fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[Import]) -> Result // Step 1) Remove the import. let remove_import_edit = autofix::edits::remove_unused_imports( - qualified_names.iter().copied(), + member_names.iter().map(AsRef::as_ref), stmt, parent, checker.locator(), @@ -217,9 +218,9 @@ fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[Import]) -> Result // Step 2) Add the import to the top-level. let add_import_edit = checker.importer().runtime_import_edit( - &StmtImports { - stmt, - qualified_names, + &ImportedMembers { + statement: stmt, + names: member_names.iter().map(AsRef::as_ref).collect(), }, at, )?; diff --git a/crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs b/crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs index 5609580ea8..5d40de8cc5 100644 --- a/crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs +++ b/crates/ruff/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs @@ -1,15 +1,17 @@ +use std::borrow::Cow; + use anyhow::Result; -use ruff_text_size::TextRange; use rustc_hash::FxHashMap; use ruff_diagnostics::{AutofixKind, Diagnostic, DiagnosticKind, Fix, Violation}; use ruff_macros::{derive_message_formats, violation}; -use ruff_python_semantic::{Binding, NodeId, ResolvedReferenceId, Scope}; +use ruff_python_semantic::{AnyImport, Binding, Imported, NodeId, ResolvedReferenceId, Scope}; +use ruff_text_size::TextRange; use crate::autofix; use crate::checkers::ast::Checker; use crate::codes::Rule; -use crate::importer::StmtImports; +use crate::importer::ImportedMembers; use crate::rules::isort::{categorize, ImportSection, ImportType}; /// ## What it does @@ -188,9 +190,9 @@ pub(crate) fn typing_only_runtime_import( diagnostics: &mut Vec, ) { // Collect all typing-only imports by statement and import type. - let mut errors_by_statement: FxHashMap<(NodeId, ImportType), Vec> = + let mut errors_by_statement: FxHashMap<(NodeId, ImportType), Vec> = FxHashMap::default(); - let mut ignores_by_statement: FxHashMap<(NodeId, ImportType), Vec> = + let mut ignores_by_statement: FxHashMap<(NodeId, ImportType), Vec> = FxHashMap::default(); for binding_id in scope.binding_ids() { @@ -206,23 +208,10 @@ pub(crate) fn typing_only_runtime_import( continue; } - let Some(qualified_name) = binding.qualified_name() else { + let Some(import) = binding.as_any_import() else { continue; }; - if is_exempt( - qualified_name, - &checker - .settings - .flake8_type_checking - .exempt_modules - .iter() - .map(String::as_str) - .collect::>(), - ) { - continue; - } - let Some(reference_id) = binding.references.first().copied() else { continue; }; @@ -236,20 +225,25 @@ pub(crate) fn typing_only_runtime_import( .is_typing() }) { - // Extract the module base and level from the full name. - // Ex) `foo.bar.baz` -> `foo`, `0` - // Ex) `.foo.bar.baz` -> `foo`, `1` - let level = qualified_name - .chars() - .take_while(|c| *c == '.') - .count() - .try_into() - .unwrap(); + let qualified_name = import.qualified_name(); + + if is_exempt( + qualified_name.as_str(), + &checker + .settings + .flake8_type_checking + .exempt_modules + .iter() + .map(String::as_str) + .collect::>(), + ) { + continue; + } // Categorize the import, using coarse-grained categorization. let import_type = match categorize( - qualified_name, - Some(level), + qualified_name.as_str(), + None, &checker.settings.src, checker.package(), &checker.settings.isort.known_modules, @@ -275,8 +269,8 @@ pub(crate) fn typing_only_runtime_import( continue; }; - let import = Import { - qualified_name, + let import = ImportBinding { + import, reference_id, range: binding.range, parent_range: binding.parent_range(checker.semantic()), @@ -309,17 +303,15 @@ pub(crate) fn typing_only_runtime_import( None }; - for Import { - qualified_name, + for ImportBinding { + import, range, parent_range, .. } in imports { - let mut diagnostic = Diagnostic::new( - diagnostic_for(import_type, qualified_name.to_string()), - range, - ); + let mut diagnostic = + Diagnostic::new(diagnostic_for(import_type, import.qualified_name()), range); if let Some(range) = parent_range { diagnostic.set_parent(range.start()); } @@ -333,17 +325,15 @@ pub(crate) fn typing_only_runtime_import( // Separately, generate a diagnostic for every _ignored_ import, to ensure that the // suppression comments aren't marked as unused. for ((_, import_type), imports) in ignores_by_statement { - for Import { - qualified_name, + for ImportBinding { + import, range, parent_range, .. } in imports { - let mut diagnostic = Diagnostic::new( - diagnostic_for(import_type, qualified_name.to_string()), - range, - ); + let mut diagnostic = + Diagnostic::new(diagnostic_for(import_type, import.qualified_name()), range); if let Some(range) = parent_range { diagnostic.set_parent(range.start()); } @@ -353,9 +343,9 @@ pub(crate) fn typing_only_runtime_import( } /// A runtime-required import with its surrounding context. -struct Import<'a> { +struct ImportBinding<'a> { /// The qualified name of the import (e.g., `typing.List` for `from typing import List`). - qualified_name: &'a str, + import: AnyImport<'a>, /// The first reference to the imported symbol. reference_id: ResolvedReferenceId, /// The trimmed range of the import (e.g., `List` in `from typing import List`). @@ -386,13 +376,13 @@ fn diagnostic_for(import_type: ImportType, qualified_name: String) -> Diagnostic /// Return `true` if `this` is implicitly loaded via importing `that`. fn is_implicit_import(this: &Binding, that: &Binding) -> bool { - let Some(this_module) = this.module_name() else { + let Some(this_import) = this.as_any_import() else { return false; }; - let Some(that_module) = that.module_name() else { + let Some(that_import) = that.as_any_import() else { return false; }; - this_module == that_module + this_import.module_name() == that_import.module_name() } /// Return `true` if `name` is exempt from typing-only enforcement. @@ -412,18 +402,18 @@ fn is_exempt(name: &str, exempt_modules: &[&str]) -> bool { } /// Generate a [`Fix`] to remove typing-only imports from a runtime context. -fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[Import]) -> Result { +fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[ImportBinding]) -> Result { let stmt = checker.semantic().stmts[stmt_id]; let parent = checker.semantic().stmts.parent(stmt); - let qualified_names: Vec<&str> = imports + let member_names: Vec> = imports .iter() - .map(|Import { qualified_name, .. }| *qualified_name) + .map(|ImportBinding { import, .. }| import.member_name()) .collect(); // Find the first reference across all imports. let at = imports .iter() - .map(|Import { reference_id, .. }| { + .map(|ImportBinding { reference_id, .. }| { checker.semantic().reference(*reference_id).range().start() }) .min() @@ -431,7 +421,7 @@ fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[Import]) -> Result // Step 1) Remove the import. let remove_import_edit = autofix::edits::remove_unused_imports( - qualified_names.iter().copied(), + member_names.iter().map(AsRef::as_ref), stmt, parent, checker.locator(), @@ -441,9 +431,9 @@ fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[Import]) -> Result // Step 2) Add the import to a `TYPE_CHECKING` block. let add_import_edit = checker.importer().typing_import_edit( - &StmtImports { - stmt, - qualified_names, + &ImportedMembers { + statement: stmt, + names: member_names.iter().map(AsRef::as_ref).collect(), }, at, checker.semantic(), diff --git a/crates/ruff/src/rules/pandas_vet/helpers.rs b/crates/ruff/src/rules/pandas_vet/helpers.rs index 9da577d606..1136bfa1c7 100644 --- a/crates/ruff/src/rules/pandas_vet/helpers.rs +++ b/crates/ruff/src/rules/pandas_vet/helpers.rs @@ -1,8 +1,8 @@ use ruff_python_ast as ast; use ruff_python_ast::Expr; +use ruff_python_semantic::{BindingKind, Imported, SemanticModel}; -use ruff_python_semantic::{BindingKind, Import, SemanticModel}; - +#[derive(Debug)] pub(super) enum Resolution { /// The expression resolves to an irrelevant expression type (e.g., a constant). IrrelevantExpression, @@ -26,26 +26,23 @@ pub(super) fn test_expression(expr: &Expr, semantic: &SemanticModel) -> Resoluti | Expr::ListComp(_) | Expr::DictComp(_) | Expr::GeneratorExp(_) => Resolution::IrrelevantExpression, - Expr::Name(ast::ExprName { id, .. }) => { - semantic - .find_binding(id) - .map_or(Resolution::IrrelevantBinding, |binding| { - match binding.kind { - BindingKind::Annotation - | BindingKind::Argument - | BindingKind::Assignment - | BindingKind::NamedExprAssignment - | BindingKind::UnpackedAssignment - | BindingKind::LoopVar - | BindingKind::Global - | BindingKind::Nonlocal(_) => Resolution::RelevantLocal, - BindingKind::Import(Import { - qualified_name: module, - }) if module == "pandas" => Resolution::PandasModule, - _ => Resolution::IrrelevantBinding, - } - }) - } + Expr::Name(ast::ExprName { id, .. }) => semantic.find_binding(id).map_or( + Resolution::IrrelevantBinding, + |binding| match &binding.kind { + BindingKind::Annotation + | BindingKind::Argument + | BindingKind::Assignment + | BindingKind::NamedExprAssignment + | BindingKind::UnpackedAssignment + | BindingKind::LoopVar + | BindingKind::Global + | BindingKind::Nonlocal(_) => Resolution::RelevantLocal, + BindingKind::Import(import) if matches!(import.call_path(), ["pandas"]) => { + Resolution::PandasModule + } + _ => Resolution::IrrelevantBinding, + }, + ), _ => Resolution::RelevantLocal, } } diff --git a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs index 71cb125d7b..67f27c51f7 100644 --- a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs +++ b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs @@ -2,7 +2,8 @@ use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation}; use ruff_macros::{derive_message_formats, violation}; use ruff_python_ast::helpers::is_const_true; use ruff_python_ast::{self as ast, Keyword, PySourceType, Ranged}; -use ruff_python_semantic::{BindingKind, Import}; +use ruff_python_semantic::BindingKind; +use ruff_python_semantic::Imported; use ruff_source_file::Locator; use crate::autofix::edits::{remove_argument, Parentheses}; @@ -57,12 +58,11 @@ pub(crate) fn inplace_argument(checker: &mut Checker, call: &ast::ExprCall) { .first() .and_then(|module| checker.semantic().find_binding(module)) .is_some_and(|binding| { - matches!( - binding.kind, - BindingKind::Import(Import { - qualified_name: "pandas" - }) - ) + if let BindingKind::Import(import) = &binding.kind { + matches!(import.call_path(), ["pandas"]) + } else { + false + } }) { return; @@ -94,10 +94,10 @@ pub(crate) fn inplace_argument(checker: &mut Checker, call: &ast::ExprCall) { && !checker.semantic().scope().kind.is_lambda() { if let Some(fix) = convert_inplace_argument_to_assignment( - checker.locator(), call, keyword, checker.source_type, + checker.locator(), ) { diagnostic.set_fix(fix); } @@ -116,10 +116,10 @@ pub(crate) fn inplace_argument(checker: &mut Checker, call: &ast::ExprCall) { /// Remove the `inplace` argument from a function call and replace it with an /// assignment. fn convert_inplace_argument_to_assignment( - locator: &Locator, call: &ast::ExprCall, keyword: &Keyword, source_type: PySourceType, + locator: &Locator, ) -> Option { // Add the assignment. let attr = call.func.as_attribute_expr()?; diff --git a/crates/ruff/src/rules/pyflakes/rules/unused_import.rs b/crates/ruff/src/rules/pyflakes/rules/unused_import.rs index 01e4ead7ea..54aaea3f40 100644 --- a/crates/ruff/src/rules/pyflakes/rules/unused_import.rs +++ b/crates/ruff/src/rules/pyflakes/rules/unused_import.rs @@ -1,10 +1,11 @@ use anyhow::Result; -use ruff_text_size::TextRange; use rustc_hash::FxHashMap; +use std::borrow::Cow; use ruff_diagnostics::{AutofixKind, Diagnostic, Fix, Violation}; use ruff_macros::{derive_message_formats, violation}; -use ruff_python_semantic::{Exceptions, NodeId, Scope}; +use ruff_python_semantic::{AnyImport, Exceptions, Imported, NodeId, Scope}; +use ruff_text_size::TextRange; use crate::autofix; use crate::checkers::ast::Checker; @@ -97,8 +98,8 @@ impl Violation for UnusedImport { pub(crate) fn unused_import(checker: &Checker, scope: &Scope, diagnostics: &mut Vec) { // Collect all unused imports by statement. - let mut unused: FxHashMap<(NodeId, Exceptions), Vec> = FxHashMap::default(); - let mut ignored: FxHashMap<(NodeId, Exceptions), Vec> = FxHashMap::default(); + let mut unused: FxHashMap<(NodeId, Exceptions), Vec> = FxHashMap::default(); + let mut ignored: FxHashMap<(NodeId, Exceptions), Vec> = FxHashMap::default(); for binding_id in scope.binding_ids() { let binding = checker.semantic().binding(binding_id); @@ -111,7 +112,7 @@ pub(crate) fn unused_import(checker: &Checker, scope: &Scope, diagnostics: &mut continue; } - let Some(qualified_name) = binding.qualified_name() else { + let Some(import) = binding.as_any_import() else { continue; }; @@ -119,8 +120,8 @@ pub(crate) fn unused_import(checker: &Checker, scope: &Scope, diagnostics: &mut continue; }; - let import = Import { - qualified_name, + let import = ImportBinding { + import, range: binding.range, parent_range: binding.parent_range(checker.semantic()), }; @@ -158,15 +159,15 @@ pub(crate) fn unused_import(checker: &Checker, scope: &Scope, diagnostics: &mut None }; - for Import { - qualified_name, + for ImportBinding { + import, range, parent_range, } in imports { let mut diagnostic = Diagnostic::new( UnusedImport { - name: qualified_name.to_string(), + name: import.qualified_name(), context: if in_except_handler { Some(UnusedImportContext::ExceptHandler) } else if in_init { @@ -192,15 +193,15 @@ pub(crate) fn unused_import(checker: &Checker, scope: &Scope, diagnostics: &mut // Separately, generate a diagnostic for every _ignored_ import, to ensure that the // suppression comments aren't marked as unused. - for Import { - qualified_name, + for ImportBinding { + import, range, parent_range, } in ignored.into_values().flatten() { let mut diagnostic = Diagnostic::new( UnusedImport { - name: qualified_name.to_string(), + name: import.qualified_name(), context: None, multiple: false, }, @@ -214,9 +215,9 @@ pub(crate) fn unused_import(checker: &Checker, scope: &Scope, diagnostics: &mut } /// An unused import with its surrounding context. -struct Import<'a> { +struct ImportBinding<'a> { /// The qualified name of the import (e.g., `typing.List` for `from typing import List`). - qualified_name: &'a str, + import: AnyImport<'a>, /// The trimmed range of the import (e.g., `List` in `from typing import List`). range: TextRange, /// The range of the import's parent statement. @@ -224,13 +225,17 @@ struct Import<'a> { } /// Generate a [`Fix`] to remove unused imports from a statement. -fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[Import]) -> Result { +fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[ImportBinding]) -> Result { let stmt = checker.semantic().stmts[stmt_id]; let parent = checker.semantic().stmts.parent(stmt); + + let member_names: Vec> = imports + .iter() + .map(|ImportBinding { import, .. }| import.member_name()) + .collect(); + let edit = autofix::edits::remove_unused_imports( - imports - .iter() - .map(|Import { qualified_name, .. }| *qualified_name), + member_names.iter().map(AsRef::as_ref), stmt, parent, checker.locator(), diff --git a/crates/ruff/src/rules/pyupgrade/rules/f_strings.rs b/crates/ruff/src/rules/pyupgrade/rules/f_strings.rs index 37dcabe00b..a4922f36a2 100644 --- a/crates/ruff/src/rules/pyupgrade/rules/f_strings.rs +++ b/crates/ruff/src/rules/pyupgrade/rules/f_strings.rs @@ -318,7 +318,10 @@ pub(crate) fn f_strings( return; } - let Expr::Call(ast::ExprCall { func, arguments, .. }) = expr else { + let Expr::Call(ast::ExprCall { + func, arguments, .. + }) = expr + else { return; }; diff --git a/crates/ruff/src/rules/pyupgrade/rules/os_error_alias.rs b/crates/ruff/src/rules/pyupgrade/rules/os_error_alias.rs index 8898c4526c..9ec667999d 100644 --- a/crates/ruff/src/rules/pyupgrade/rules/os_error_alias.rs +++ b/crates/ruff/src/rules/pyupgrade/rules/os_error_alias.rs @@ -36,7 +36,7 @@ use crate::registry::AsRule; /// - [Python documentation: `OSError`](https://docs.python.org/3/library/exceptions.html#OSError) #[violation] pub struct OSErrorAlias { - pub name: Option, + name: Option, } impl AlwaysAutofixableViolation for OSErrorAlias { diff --git a/crates/ruff/src/rules/pyupgrade/rules/unnecessary_builtin_import.rs b/crates/ruff/src/rules/pyupgrade/rules/unnecessary_builtin_import.rs index 01e12717a8..17fbcbe284 100644 --- a/crates/ruff/src/rules/pyupgrade/rules/unnecessary_builtin_import.rs +++ b/crates/ruff/src/rules/pyupgrade/rules/unnecessary_builtin_import.rs @@ -124,12 +124,11 @@ pub(crate) fn unnecessary_builtin_import( diagnostic.try_set_fix(|| { let stmt = checker.semantic().stmt(); let parent = checker.semantic().stmt_parent(); - let unused_imports: Vec = unused_imports - .iter() - .map(|alias| format!("{module}.{}", alias.name)) - .collect(); let edit = autofix::edits::remove_unused_imports( - unused_imports.iter().map(String::as_str), + unused_imports + .iter() + .map(|alias| &alias.name) + .map(ruff_python_ast::Identifier::as_str), stmt, parent, checker.locator(), diff --git a/crates/ruff/src/rules/pyupgrade/rules/unnecessary_future_import.rs b/crates/ruff/src/rules/pyupgrade/rules/unnecessary_future_import.rs index 52c5b543c8..ec4e4d7ba8 100644 --- a/crates/ruff/src/rules/pyupgrade/rules/unnecessary_future_import.rs +++ b/crates/ruff/src/rules/pyupgrade/rules/unnecessary_future_import.rs @@ -111,14 +111,13 @@ pub(crate) fn unnecessary_future_import(checker: &mut Checker, stmt: &Stmt, name if checker.patch(diagnostic.kind.rule()) { diagnostic.try_set_fix(|| { - let unused_imports: Vec = unused_imports - .iter() - .map(|alias| format!("__future__.{}", alias.name)) - .collect(); let stmt = checker.semantic().stmt(); let parent = checker.semantic().stmt_parent(); let edit = autofix::edits::remove_unused_imports( - unused_imports.iter().map(String::as_str), + unused_imports + .iter() + .map(|alias| &alias.name) + .map(ruff_python_ast::Identifier::as_str), stmt, parent, checker.locator(), diff --git a/crates/ruff_python_ast/src/call_path.rs b/crates/ruff_python_ast/src/call_path.rs index c11c456fe7..fe8043a106 100644 --- a/crates/ruff_python_ast/src/call_path.rs +++ b/crates/ruff_python_ast/src/call_path.rs @@ -1,6 +1,7 @@ -use crate::{nodes, Expr}; use smallvec::{smallvec, SmallVec}; +use crate::{nodes, Expr}; + /// A representation of a qualified name, like `typing.List`. pub type CallPath<'a> = SmallVec<[&'a str; 8]>; @@ -141,12 +142,31 @@ pub fn compose_call_path(expr: &Expr) -> Option { /// Format a call path for display. pub fn format_call_path(call_path: &[&str]) -> String { - if call_path - .first() - .expect("Unable to format empty call path") - .is_empty() - { + if call_path.first().map_or(false, |first| first.is_empty()) { + // If the first segment is empty, the `CallPath` is that of a builtin. + // Ex) `["", "bool"]` -> `"bool"` call_path[1..].join(".") + } else if call_path + .first() + .map_or(false, |first| matches!(*first, ".")) + { + // If the call path is dot-prefixed, it's an unresolved relative import. + // Ex) `[".foo", "bar"]` -> `".foo.bar"` + let mut formatted = String::new(); + let mut iter = call_path.iter(); + for segment in iter.by_ref() { + if *segment == "." { + formatted.push('.'); + } else { + formatted.push_str(segment); + break; + } + } + for segment in iter { + formatted.push('.'); + formatted.push_str(segment); + } + formatted } else { call_path.join(".") } diff --git a/crates/ruff_python_ast/src/helpers.rs b/crates/ruff_python_ast/src/helpers.rs index 2666f32de3..d29a68b167 100644 --- a/crates/ruff_python_ast/src/helpers.rs +++ b/crates/ruff_python_ast/src/helpers.rs @@ -790,42 +790,81 @@ pub fn to_module_path(package: &Path, path: &Path) -> Option> { .collect::>>() } -/// 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). +/// Format the call path for a relative import. /// /// # Examples /// /// ```rust -/// # use smallvec::{smallvec, SmallVec}; -/// # use ruff_python_ast::helpers::from_relative_import; +/// # use ruff_python_ast::helpers::collect_import_from_member; /// -/// 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()); +/// assert_eq!(collect_import_from_member(None, None, "bar").as_slice(), ["bar"]); +/// assert_eq!(collect_import_from_member(Some(1), None, "bar").as_slice(), [".", "bar"]); +/// assert_eq!(collect_import_from_member(Some(1), Some("foo"), "bar").as_slice(), [".", "foo", "bar"]); /// ``` -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); +pub fn collect_import_from_member<'a>( + level: Option, + module: Option<&'a str>, + member: &'a str, +) -> CallPath<'a> { + let mut call_path: CallPath = SmallVec::with_capacity( + level.unwrap_or_default() as usize + + module + .map(|module| module.split('.').count()) + .unwrap_or_default() + + 1, + ); + + // Include the dots as standalone segments. + if let Some(level) = level { + if level > 0 { + for _ in 0..level { + call_path.push("."); + } + } + } + + // Add the remaining segments. + if let Some(module) = module { + call_path.extend(module.split('.')); + } + + // Add the member. + call_path.push(member); + + call_path +} + +/// Format the call path for a relative import, or `None` if the relative import extends beyond +/// the root module. +pub fn from_relative_import<'a>( + // The path from which the import is relative. + module: &'a [String], + // The path of the import itself (e.g., given `from ..foo import bar`, `[".", ".", "foo", "bar]`). + import: &[&'a str], + // The remaining segments to the call path (e.g., given `bar.baz`, `["baz"]`). + tail: &[&'a str], +) -> Option> { + let mut call_path: CallPath = SmallVec::with_capacity(module.len() + import.len() + tail.len()); // Start with the module path. call_path.extend(module.iter().map(String::as_str)); // 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(); + for segment in import { + if *segment == "." { + if call_path.is_empty() { + return None; + } + call_path.pop(); + } else { + call_path.push(segment); } - call_path.pop(); } // Add the remaining segments. - call_path.extend(name.trim_start_matches('.').split('.')); + call_path.extend_from_slice(tail); - call_path + Some(call_path) } /// Given an imported module (based on its relative import level and module name), return the diff --git a/crates/ruff_python_semantic/src/binding.rs b/crates/ruff_python_semantic/src/binding.rs index 7dd9056dd6..10e1900888 100644 --- a/crates/ruff_python_semantic/src/binding.rs +++ b/crates/ruff_python_semantic/src/binding.rs @@ -1,11 +1,13 @@ +use std::borrow::Cow; use std::ops::{Deref, DerefMut}; use bitflags::bitflags; -use ruff_python_ast::Ranged; -use ruff_text_size::TextRange; use ruff_index::{newtype_index, IndexSlice, IndexVec}; +use ruff_python_ast::call_path::format_call_path; +use ruff_python_ast::Ranged; use ruff_source_file::Locator; +use ruff_text_size::TextRange; use crate::context::ExecutionContext; use crate::model::SemanticModel; @@ -117,38 +119,38 @@ impl<'a> Binding<'a> { // import foo.baz // ``` BindingKind::Import(Import { - qualified_name: redefinition, + call_path: redefinition, }) => { if let BindingKind::SubmoduleImport(SubmoduleImport { - qualified_name: definition, + call_path: definition, }) = &existing.kind { return redefinition == definition; } } BindingKind::FromImport(FromImport { - qualified_name: redefinition, + call_path: redefinition, }) => { if let BindingKind::SubmoduleImport(SubmoduleImport { - qualified_name: definition, + call_path: definition, }) = &existing.kind { return redefinition == definition; } } BindingKind::SubmoduleImport(SubmoduleImport { - qualified_name: redefinition, + call_path: redefinition, }) => match &existing.kind { BindingKind::Import(Import { - qualified_name: definition, + call_path: definition, }) | BindingKind::SubmoduleImport(SubmoduleImport { - qualified_name: definition, + call_path: definition, }) => { return redefinition == definition; } BindingKind::FromImport(FromImport { - qualified_name: definition, + call_path: definition, }) => { return redefinition == definition; } @@ -175,35 +177,6 @@ impl<'a> Binding<'a> { ) } - /// Returns the fully-qualified symbol name, if this symbol was imported from another module. - pub fn qualified_name(&self) -> Option<&str> { - match &self.kind { - BindingKind::Import(Import { qualified_name }) => Some(qualified_name), - BindingKind::FromImport(FromImport { qualified_name }) => Some(qualified_name), - BindingKind::SubmoduleImport(SubmoduleImport { qualified_name }) => { - Some(qualified_name) - } - _ => None, - } - } - - /// Returns the fully-qualified name of the module from which this symbol was imported, if this - /// symbol was imported from another module. - pub fn module_name(&self) -> Option<&str> { - match &self.kind { - BindingKind::Import(Import { qualified_name }) - | BindingKind::SubmoduleImport(SubmoduleImport { qualified_name }) => { - Some(qualified_name.split('.').next().unwrap_or(qualified_name)) - } - BindingKind::FromImport(FromImport { qualified_name }) => Some( - qualified_name - .rsplit_once('.') - .map_or(qualified_name, |(module, _)| module), - ), - _ => None, - } - } - /// Returns the name of the binding (e.g., `x` in `x = 1`). pub fn name<'b>(&self, locator: &'b Locator) -> &'b str { locator.slice(self.range) @@ -221,6 +194,15 @@ impl<'a> Binding<'a> { } }) } + + pub fn as_any_import(&'a self) -> Option> { + match &self.kind { + BindingKind::Import(import) => Some(AnyImport::Import(import)), + BindingKind::SubmoduleImport(import) => Some(AnyImport::SubmoduleImport(import)), + BindingKind::FromImport(import) => Some(AnyImport::FromImport(import)), + _ => None, + } + } } bitflags! { @@ -356,18 +338,18 @@ pub struct Import<'a> { /// The full name of the module being imported. /// Ex) Given `import foo`, `qualified_name` would be "foo". /// Ex) Given `import foo as bar`, `qualified_name` would be "foo". - pub qualified_name: &'a str, + pub call_path: Box<[&'a str]>, } /// A binding for a member imported from a module, keyed on the name to which the member is bound. /// Ex) `from foo import bar` would be keyed on "bar". /// Ex) `from foo import bar as baz` would be keyed on "baz". #[derive(Debug, Clone)] -pub struct FromImport { +pub struct FromImport<'a> { /// The full name of the member being imported. /// Ex) Given `from foo import bar`, `qualified_name` would be "foo.bar". /// Ex) Given `from foo import bar as baz`, `qualified_name` would be "foo.bar". - pub qualified_name: String, + pub call_path: Box<[&'a str]>, } /// A binding for a submodule imported from a module, keyed on the name of the parent module. @@ -376,7 +358,7 @@ pub struct FromImport { pub struct SubmoduleImport<'a> { /// The full name of the submodule being imported. /// Ex) Given `import foo.bar`, `qualified_name` would be "foo.bar". - pub qualified_name: &'a str, + pub call_path: Box<[&'a str]>, } #[derive(Debug, Clone, is_macro::Is)] @@ -485,7 +467,7 @@ pub enum BindingKind<'a> { /// ```python /// from foo import bar /// ``` - FromImport(FromImport), + FromImport(FromImport<'a>), /// A binding for a submodule imported from a module, like `bar` in: /// ```python @@ -532,3 +514,106 @@ bitflags! { const IMPORT_ERROR = 0b0000_0100; } } + +/// A trait for imported symbols. +pub trait Imported<'a> { + /// Returns the call path to the imported symbol. + fn call_path(&self) -> &[&str]; + + /// Returns the module name of the imported symbol. + fn module_name(&self) -> &[&str]; + + /// Returns the member name of the imported symbol. For a straight import, this is equivalent + /// to the qualified name; for a `from` import, this is the name of the imported symbol. + fn member_name(&self) -> Cow<'a, str>; + + /// Returns the fully-qualified name of the imported symbol. + fn qualified_name(&self) -> String { + format_call_path(self.call_path()) + } +} + +impl<'a> Imported<'a> for Import<'a> { + /// For example, given `import foo`, returns `["foo"]`. + fn call_path(&self) -> &[&str] { + self.call_path.as_ref() + } + + /// For example, given `import foo`, returns `["foo"]`. + fn module_name(&self) -> &[&str] { + &self.call_path[..1] + } + + /// For example, given `import foo`, returns `"foo"`. + fn member_name(&self) -> Cow<'a, str> { + Cow::Owned(self.qualified_name()) + } +} + +impl<'a> Imported<'a> for SubmoduleImport<'a> { + /// For example, given `import foo.bar`, returns `["foo", "bar"]`. + fn call_path(&self) -> &[&str] { + self.call_path.as_ref() + } + + /// For example, given `import foo.bar`, returns `["foo"]`. + fn module_name(&self) -> &[&str] { + &self.call_path[..1] + } + + /// For example, given `import foo.bar`, returns `"foo.bar"`. + fn member_name(&self) -> Cow<'a, str> { + Cow::Owned(self.qualified_name()) + } +} + +impl<'a> Imported<'a> for FromImport<'a> { + /// For example, given `from foo import bar`, returns `["foo", "bar"]`. + fn call_path(&self) -> &[&str] { + self.call_path.as_ref() + } + + /// For example, given `from foo import bar`, returns `["foo"]`. + fn module_name(&self) -> &[&str] { + &self.call_path[..self.call_path.len() - 1] + } + + /// For example, given `from foo import bar`, returns `"bar"`. + fn member_name(&self) -> Cow<'a, str> { + Cow::Borrowed(self.call_path[self.call_path.len() - 1]) + } +} + +/// A wrapper around an import [`BindingKind`] that can be any of the three types of imports. +#[derive(Debug, Clone)] +pub enum AnyImport<'a> { + Import(&'a Import<'a>), + SubmoduleImport(&'a SubmoduleImport<'a>), + FromImport(&'a FromImport<'a>), +} + +impl<'a> Imported<'a> for AnyImport<'a> { + fn call_path(&self) -> &[&str] { + match self { + Self::Import(import) => import.call_path(), + Self::SubmoduleImport(import) => import.call_path(), + Self::FromImport(import) => import.call_path(), + } + } + + fn module_name(&self) -> &[&str] { + match self { + Self::Import(import) => import.module_name(), + Self::SubmoduleImport(import) => import.module_name(), + Self::FromImport(import) => import.module_name(), + } + } + + fn member_name(&self) -> Cow<'a, str> { + match self { + Self::Import(import) => import.member_name(), + Self::SubmoduleImport(import) => import.member_name(), + Self::FromImport(import) => import.member_name(), + } + } +} diff --git a/crates/ruff_python_semantic/src/model.rs b/crates/ruff_python_semantic/src/model.rs index 3f9711bd84..bede85c5fd 100644 --- a/crates/ruff_python_semantic/src/model.rs +++ b/crates/ruff_python_semantic/src/model.rs @@ -1,15 +1,15 @@ use std::path::Path; use bitflags::bitflags; -use ruff_python_ast::{self as ast, Expr, Ranged, Stmt}; -use ruff_text_size::{TextRange, TextSize}; use rustc_hash::FxHashMap; use smallvec::smallvec; use ruff_python_ast::call_path::{collect_call_path, from_unqualified_name, CallPath}; use ruff_python_ast::helpers::from_relative_import; +use ruff_python_ast::{self as ast, Expr, Ranged, Stmt}; use ruff_python_stdlib::path::is_python_stub_file; use ruff_python_stdlib::typing::is_typing_extension; +use ruff_text_size::{TextRange, TextSize}; use crate::binding::{ Binding, BindingFlags, BindingId, BindingKind, Bindings, Exceptions, FromImport, Import, @@ -23,7 +23,7 @@ use crate::reference::{ ResolvedReference, ResolvedReferenceId, ResolvedReferences, UnresolvedReferences, }; use crate::scope::{Scope, ScopeId, ScopeKind, Scopes}; -use crate::{UnresolvedReference, UnresolvedReferenceFlags}; +use crate::{Imported, UnresolvedReference, UnresolvedReferenceFlags}; /// A semantic model for a Python module, to enable querying the module's semantic information. pub struct SemanticModel<'a> { @@ -583,17 +583,14 @@ impl<'a> SemanticModel<'a> { // import pyarrow.csv // print(pa.csv.read_csv("test.csv")) // ``` - let qualified_name = self.bindings[binding_id].qualified_name()?; - let has_alias = qualified_name - .split('.') - .last() - .map(|segment| segment != symbol) - .unwrap_or_default(); - if !has_alias { + let import = self.bindings[binding_id].as_any_import()?; + let call_path = import.call_path(); + let segment = call_path.last()?; + if *segment == symbol { return None; } - let binding_id = self.scopes[scope_id].get(qualified_name)?; + let binding_id = self.scopes[scope_id].get(segment)?; if !self.bindings[binding_id].kind.is_submodule_import() { return None; } @@ -625,53 +622,41 @@ impl<'a> SemanticModel<'a> { // If the name was already resolved, look it up; otherwise, search for the symbol. let head = match_head(value)?; - let binding = self - .resolved_names - .get(&head.into()) - .map(|id| self.binding(*id)) - .or_else(|| self.find_binding(&head.id))?; + let binding = if let Some(id) = self.resolved_names.get(&head.into()) { + self.binding(*id) + } else { + self.find_binding(&head.id)? + }; match &binding.kind { - BindingKind::Import(Import { - qualified_name: name, - }) => { - let call_path = collect_call_path(value)?; - let (_, tail) = call_path.split_first()?; - - let mut source_path: CallPath = from_unqualified_name(name); - source_path.extend_from_slice(tail); - Some(source_path) + BindingKind::Import(Import { call_path }) => { + let value_path = collect_call_path(value)?; + let (_, tail) = value_path.split_first()?; + let resolved: CallPath = call_path.iter().chain(tail.iter()).copied().collect(); + Some(resolved) } - BindingKind::SubmoduleImport(SubmoduleImport { - qualified_name: name, - }) => { - let call_path = collect_call_path(value)?; - let (_, tail) = call_path.split_first()?; - - let name = name.split('.').next().unwrap_or(name); - let mut source_path: CallPath = from_unqualified_name(name); - source_path.extend_from_slice(tail); - Some(source_path) + BindingKind::SubmoduleImport(SubmoduleImport { call_path }) => { + let value_path = collect_call_path(value)?; + let (_, tail) = value_path.split_first()?; + let resolved: CallPath = call_path + .iter() + .take(1) + .chain(tail.iter()) + .copied() + .collect(); + Some(resolved) } - BindingKind::FromImport(FromImport { - qualified_name: name, - }) => { - let call_path = collect_call_path(value)?; - let (_, tail) = call_path.split_first()?; + BindingKind::FromImport(FromImport { call_path }) => { + let value_path = collect_call_path(value)?; + let (_, tail) = value_path.split_first()?; - if name.starts_with('.') { - let mut source_path = from_relative_import(self.module_path?, name); - if source_path.is_empty() { - None + let resolved: CallPath = + if call_path.first().map_or(false, |segment| *segment == ".") { + from_relative_import(self.module_path?, call_path, tail)? } else { - source_path.extend_from_slice(tail); - Some(source_path) - } - } else { - let mut source_path: CallPath = from_unqualified_name(name); - source_path.extend_from_slice(tail); - Some(source_path) - } + call_path.iter().chain(tail.iter()).copied().collect() + }; + Some(resolved) } BindingKind::Builtin => Some(smallvec!["", head.id.as_str()]), _ => None, @@ -695,6 +680,8 @@ impl<'a> SemanticModel<'a> { module: &str, member: &str, ) -> Option { + // TODO(charlie): Pass in a slice. + let module_path: Vec<&str> = module.split('.').collect(); self.scopes().enumerate().find_map(|(scope_index, scope)| { scope.bindings().find_map(|(name, binding_id)| { let binding = &self.bindings[binding_id]; @@ -702,8 +689,8 @@ impl<'a> SemanticModel<'a> { // Ex) Given `module="sys"` and `object="exit"`: // `import sys` -> `sys.exit` // `import sys as sys2` -> `sys2.exit` - BindingKind::Import(Import { qualified_name }) => { - if qualified_name == &module { + BindingKind::Import(Import { call_path }) => { + if call_path.as_ref() == module_path.as_slice() { if let Some(source) = binding.source { // Verify that `sys` isn't bound in an inner scope. if self @@ -723,10 +710,9 @@ impl<'a> SemanticModel<'a> { // Ex) Given `module="os.path"` and `object="join"`: // `from os.path import join` -> `join` // `from os.path import join as join2` -> `join2` - BindingKind::FromImport(FromImport { qualified_name }) => { - if let Some((target_module, target_member)) = qualified_name.split_once('.') - { - if target_module == module && target_member == member { + BindingKind::FromImport(FromImport { call_path }) => { + if let Some((target_member, target_module)) = call_path.split_last() { + if target_module == module_path.as_slice() && target_member == &member { if let Some(source) = binding.source { // Verify that `join` isn't bound in an inner scope. if self