From ac25bd9596f3eb6a5e05a617364bb66d1f3ed6fd Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Thu, 27 Jul 2023 21:08:08 -0400 Subject: [PATCH] Traits --- crates/ruff/src/checkers/ast/mod.rs | 11 +- .../rules/unconventional_import_alias.rs | 8 +- .../unaliased_collections_abc_set_import.rs | 7 +- .../runtime_import_in_type_checking_block.rs | 2 +- .../rules/typing_only_runtime_import.rs | 8 +- crates/ruff/src/rules/pandas_vet/helpers.rs | 8 +- .../pandas_vet/rules/inplace_argument.rs | 11 +- .../src/rules/pyflakes/rules/unused_import.rs | 2 +- ...ules__pyflakes__tests__F401_F401_6.py.snap | 8 +- crates/ruff_python_ast/src/call_path.rs | 138 +------------ crates/ruff_python_ast/src/helpers.rs | 107 ++++------ crates/ruff_python_semantic/src/binding.rs | 192 ++++++++++-------- crates/ruff_python_semantic/src/model.rs | 39 +++- 13 files changed, 202 insertions(+), 339 deletions(-) diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs index b6205386b0..efd91ca4b9 100644 --- a/crates/ruff/src/checkers/ast/mod.rs +++ b/crates/ruff/src/checkers/ast/mod.rs @@ -39,7 +39,7 @@ 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, from_relative_import_parts, literal_path, to_module_path, + collect_import_from_member, extract_handled_exceptions, to_module_path, }; use ruff_python_ast::identifier::Identifier; use ruff_python_ast::str::trailing_quote; @@ -394,13 +394,8 @@ where // 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 = self - .module_path - .and_then(|module_path| { - from_relative_import_parts(module_path, level, module, &alias.name) - }) - .unwrap_or_else(|| literal_path(level, module, &alias.name)); - let call_path: Box<[&str]> = call_path.into_boxed_slice(); + let call_path = collect_import_from_member(level, module, &alias.name) + .into_boxed_slice(); self.add_binding( name, alias.identifier(), 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 528166c9d0..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,10 +56,12 @@ 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 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 581baefac2..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 { call_path }) = &binding.kind else { + let BindingKind::FromImport(import) = &binding.kind else { return None; }; - if !matches!(call_path.as_ref(), ["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 ccbb3674e1..3d867f70e8 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 @@ -5,7 +5,7 @@ use std::borrow::Cow; use ruff_diagnostics::{AutofixKind, Diagnostic, Fix, Violation}; use ruff_macros::{derive_message_formats, violation}; -use ruff_python_semantic::{AnyImport, NodeId, ResolvedReferenceId, Scope}; +use ruff_python_semantic::{AnyImport, Imported, NodeId, ResolvedReferenceId, Scope}; use crate::autofix; use crate::checkers::ast::Checker; 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 19eb915271..3515d8c83f 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 @@ -5,7 +5,7 @@ use rustc_hash::FxHashMap; use ruff_diagnostics::{AutofixKind, Diagnostic, DiagnosticKind, Fix, Violation}; use ruff_macros::{derive_message_formats, violation}; -use ruff_python_semantic::{AnyImport, Binding, NodeId, ResolvedReferenceId, Scope}; +use ruff_python_semantic::{AnyImport, Binding, Imported, NodeId, ResolvedReferenceId, Scope}; use ruff_text_size::TextRange; use crate::autofix; @@ -376,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. diff --git a/crates/ruff/src/rules/pandas_vet/helpers.rs b/crates/ruff/src/rules/pandas_vet/helpers.rs index 1aa1f0ab13..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, @@ -37,9 +37,7 @@ pub(super) fn test_expression(expr: &Expr, semantic: &SemanticModel) -> Resoluti | BindingKind::LoopVar | BindingKind::Global | BindingKind::Nonlocal(_) => Resolution::RelevantLocal, - BindingKind::Import(Import { call_path }) - if matches!(call_path.as_ref(), ["pandas"]) => - { + BindingKind::Import(import) if matches!(import.call_path(), ["pandas"]) => { Resolution::PandasModule } _ => Resolution::IrrelevantBinding, 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 ab1bf3d870..7383eecf2d 100644 --- a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs +++ b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs @@ -1,11 +1,10 @@ -use ruff_python_ast::{Expr, Keyword, Ranged}; -use ruff_text_size::TextRange; - 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_semantic::{BindingKind, Import}; +use ruff_python_ast::{Expr, Keyword, Ranged}; +use ruff_python_semantic::{BindingKind, Imported}; use ruff_source_file::Locator; +use ruff_text_size::TextRange; use crate::autofix::edits::remove_argument; use crate::checkers::ast::Checker; @@ -65,8 +64,8 @@ pub(crate) fn inplace_argument( .first() .and_then(|module| checker.semantic().find_binding(module)) .map_or(false, |binding| { - if let BindingKind::Import(Import { call_path }) = &binding.kind { - matches!(call_path.as_ref(), ["pandas"]) + if let BindingKind::Import(import) = &binding.kind { + matches!(import.call_path(), ["pandas"]) } else { false } diff --git a/crates/ruff/src/rules/pyflakes/rules/unused_import.rs b/crates/ruff/src/rules/pyflakes/rules/unused_import.rs index f03707c5bd..e13546e725 100644 --- a/crates/ruff/src/rules/pyflakes/rules/unused_import.rs +++ b/crates/ruff/src/rules/pyflakes/rules/unused_import.rs @@ -4,7 +4,7 @@ use std::borrow::Cow; use ruff_diagnostics::{AutofixKind, Diagnostic, Fix, Violation}; use ruff_macros::{derive_message_formats, violation}; -use ruff_python_semantic::{AnyImport, Exceptions, NodeId, Scope}; +use ruff_python_semantic::{AnyImport, Exceptions, Imported, NodeId, Scope}; use ruff_text_size::TextRange; use crate::autofix; diff --git a/crates/ruff/src/rules/pyflakes/snapshots/ruff__rules__pyflakes__tests__F401_F401_6.py.snap b/crates/ruff/src/rules/pyflakes/snapshots/ruff__rules__pyflakes__tests__F401_F401_6.py.snap index c5b397feea..a6b21228c3 100644 --- a/crates/ruff/src/rules/pyflakes/snapshots/ruff__rules__pyflakes__tests__F401_F401_6.py.snap +++ b/crates/ruff/src/rules/pyflakes/snapshots/ruff__rules__pyflakes__tests__F401_F401_6.py.snap @@ -1,7 +1,7 @@ --- source: crates/ruff/src/rules/pyflakes/mod.rs --- -F401_6.py:7:25: F401 [*] `pyflakes.background.BackgroundTasks` imported but unused +F401_6.py:7:25: F401 [*] `.background.BackgroundTasks` imported but unused | 6 | # F401 `background.BackgroundTasks` imported but unused 7 | from .background import BackgroundTasks @@ -9,7 +9,7 @@ F401_6.py:7:25: F401 [*] `pyflakes.background.BackgroundTasks` imported but unus 8 | 9 | # F401 `datastructures.UploadFile` imported but unused | - = help: Remove unused import: `pyflakes.background.BackgroundTasks` + = help: Remove unused import: `.background.BackgroundTasks` ℹ Fix 4 4 | from .applications import FastAPI as FastAPI @@ -20,7 +20,7 @@ F401_6.py:7:25: F401 [*] `pyflakes.background.BackgroundTasks` imported but unus 9 8 | # F401 `datastructures.UploadFile` imported but unused 10 9 | from .datastructures import UploadFile as FileUpload -F401_6.py:10:43: F401 [*] `pyflakes.datastructures.UploadFile` imported but unused +F401_6.py:10:43: F401 [*] `.datastructures.UploadFile` imported but unused | 9 | # F401 `datastructures.UploadFile` imported but unused 10 | from .datastructures import UploadFile as FileUpload @@ -28,7 +28,7 @@ F401_6.py:10:43: F401 [*] `pyflakes.datastructures.UploadFile` imported but unus 11 | 12 | # OK | - = help: Remove unused import: `pyflakes.datastructures.UploadFile` + = help: Remove unused import: `.datastructures.UploadFile` ℹ Fix 7 7 | from .background import BackgroundTasks diff --git a/crates/ruff_python_ast/src/call_path.rs b/crates/ruff_python_ast/src/call_path.rs index 72fa702c98..fe8043a106 100644 --- a/crates/ruff_python_ast/src/call_path.rs +++ b/crates/ruff_python_ast/src/call_path.rs @@ -5,142 +5,6 @@ use crate::{nodes, Expr}; /// A representation of a qualified name, like `typing.List`. pub type CallPath<'a> = SmallVec<[&'a str; 8]>; -/// Convert an `Expr` to its [`CallPath`] segments (like `["typing", "List"]`). -pub fn collect_head_path(expr: &Expr) -> Option<(&nodes::ExprName, CallPath)> { - // Unroll the loop up to eight times, to match the maximum number of expected attributes. - // In practice, unrolling appears to give about a 4x speed-up on this hot path. - let attr1 = match expr { - Expr::Attribute(attr1) => attr1, - // Ex) `foo` - Expr::Name(name) => return Some((name, CallPath::new())), - _ => return None, - }; - - let attr2 = match attr1.value.as_ref() { - Expr::Attribute(attr2) => attr2, - // Ex) `foo.bar` - Expr::Name(name) => { - return Some((name, CallPath::from_slice(&[attr1.attr.as_str()]))); - } - _ => return None, - }; - - let attr3 = match attr2.value.as_ref() { - Expr::Attribute(attr3) => attr3, - // Ex) `foo.bar.baz` - Expr::Name(name) => { - return Some(( - name, - CallPath::from_slice(&[attr2.attr.as_str(), attr1.attr.as_str()]), - )); - } - _ => return None, - }; - - let attr4 = match attr3.value.as_ref() { - Expr::Attribute(attr4) => attr4, - // Ex) `foo.bar.baz.bop` - Expr::Name(name) => { - return Some(( - name, - CallPath::from_slice(&[ - attr3.attr.as_str(), - attr2.attr.as_str(), - attr1.attr.as_str(), - ]), - )); - } - _ => return None, - }; - - let attr5 = match attr4.value.as_ref() { - Expr::Attribute(attr5) => attr5, - // Ex) `foo.bar.baz.bop.bap` - Expr::Name(name) => { - return Some(( - name, - CallPath::from_slice(&[ - attr4.attr.as_str(), - attr3.attr.as_str(), - attr2.attr.as_str(), - attr1.attr.as_str(), - ]), - )); - } - _ => return None, - }; - - let attr6 = match attr5.value.as_ref() { - Expr::Attribute(attr6) => attr6, - // Ex) `foo.bar.baz.bop.bap.bab` - Expr::Name(name) => { - return Some(( - name, - CallPath::from_slice(&[ - attr5.attr.as_str(), - attr4.attr.as_str(), - attr3.attr.as_str(), - attr2.attr.as_str(), - attr1.attr.as_str(), - ]), - )); - } - _ => return None, - }; - - let attr7 = match attr6.value.as_ref() { - Expr::Attribute(attr7) => attr7, - // Ex) `foo.bar.baz.bop.bap.bab.bob` - Expr::Name(name) => { - return Some(( - name, - CallPath::from_slice(&[ - attr6.attr.as_str(), - attr5.attr.as_str(), - attr4.attr.as_str(), - attr3.attr.as_str(), - attr2.attr.as_str(), - attr1.attr.as_str(), - ]), - )); - } - _ => return None, - }; - - let attr8 = match attr7.value.as_ref() { - Expr::Attribute(attr8) => attr8, - // Ex) `foo.bar.baz.bop.bap.bab.bob.bib` - Expr::Name(name) => { - return Some(( - name, - CallPath::from_slice(&[ - attr7.attr.as_str(), - attr6.attr.as_str(), - attr5.attr.as_str(), - attr4.attr.as_str(), - attr3.attr.as_str(), - attr2.attr.as_str(), - attr1.attr.as_str(), - ]), - )); - } - _ => return None, - }; - - let (name, mut call_path) = collect_head_path(&attr8.value)?; - call_path.extend([ - attr8.attr.as_str(), - attr7.attr.as_str(), - attr6.attr.as_str(), - attr5.attr.as_str(), - attr4.attr.as_str(), - attr3.attr.as_str(), - attr2.attr.as_str(), - attr1.attr.as_str(), - ]); - Some((name, call_path)) -} - /// Convert an `Expr` to its [`CallPath`] segments (like `["typing", "List"]`). pub fn collect_call_path(expr: &Expr) -> Option { // Unroll the loop up to eight times, to match the maximum number of expected attributes. @@ -291,7 +155,7 @@ pub fn format_call_path(call_path: &[&str]) -> String { let mut formatted = String::new(); let mut iter = call_path.iter(); for segment in iter.by_ref() { - if matches!(*segment, ".") { + if *segment == "." { formatted.push('.'); } else { formatted.push_str(segment); diff --git a/crates/ruff_python_ast/src/helpers.rs b/crates/ruff_python_ast/src/helpers.rs index 5c6683a900..5073b73f16 100644 --- a/crates/ruff_python_ast/src/helpers.rs +++ b/crates/ruff_python_ast/src/helpers.rs @@ -856,78 +856,18 @@ 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); - - // 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(); - } - call_path.pop(); - } - - // Add the remaining segments. - call_path.extend(name.trim_start_matches('.').split('.')); - - call_path -} - -pub fn from_relative_import_parts<'a>( - module_path: &'a [String], - level: Option, - module: Option<&'a str>, - member: &'a str, -) -> Option> { - let mut call_path: CallPath = SmallVec::with_capacity(module_path.len() + 1); - - // Remove segments based on the number of dots. - if let Some(level) = level { - if level > 0 { - call_path.extend(module_path.iter().map(String::as_str)); - - for _ in 0..level { - if call_path.is_empty() { - break; - } - call_path.pop(); - } - } - } - - // Add the remaining segments. - if let Some(module) = module { - call_path.extend(module.split('.')); - } - - // Add the member. - call_path.push(member); - - Some(call_path) -} - -pub fn literal_path<'a>( +pub fn collect_import_from_member<'a>( level: Option, module: Option<&'a str>, member: &'a str, @@ -940,7 +880,7 @@ pub fn literal_path<'a>( + 1, ); - // Include the dots + // Include the dots as standalone segments. if let Some(level) = level { if level > 0 { for _ in 0..level { @@ -960,6 +900,39 @@ pub fn literal_path<'a>( 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 segment in import { + if *segment == "." { + if call_path.is_empty() { + return None; + } + call_path.pop(); + } else { + call_path.push(segment); + } + } + + // Add the remaining segments. + call_path.extend_from_slice(tail); + + Some(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>( diff --git a/crates/ruff_python_semantic/src/binding.rs b/crates/ruff_python_semantic/src/binding.rs index 23c6c71221..9c48d0e052 100644 --- a/crates/ruff_python_semantic/src/binding.rs +++ b/crates/ruff_python_semantic/src/binding.rs @@ -177,50 +177,6 @@ impl<'a> Binding<'a> { ) } - /// Returns the fully-qualified symbol name, if this symbol was imported from another module. - pub fn call_path(&self) -> Option<&[&str]> { - match &self.kind { - BindingKind::Import(Import { call_path }) => Some(call_path), - BindingKind::FromImport(FromImport { call_path }) => Some(call_path), - BindingKind::SubmoduleImport(SubmoduleImport { call_path }) => Some(call_path), - _ => None, - } - } - - /// Returns the fully-qualified symbol name, if this symbol was imported from another module. - pub fn qualified_name(&self) -> Option { - self.call_path().map(format_call_path) - } - - /// 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 { call_path }) - | BindingKind::SubmoduleImport(SubmoduleImport { call_path }) => Some(&call_path[..1]), - BindingKind::FromImport(FromImport { call_path }) => { - Some(&call_path[..call_path.len() - 1]) - } - _ => None, - } - } - - /// Returns the fully-qualified symbol name, if this symbol was imported from another module. - pub fn member_name(&self) -> Option> { - match &self.kind { - BindingKind::Import(Import { call_path }) => { - Some(Cow::Owned(format_call_path(call_path))) - } - BindingKind::FromImport(FromImport { call_path }) => { - call_path.last().map(|member| Cow::Borrowed(*member)) - } - BindingKind::SubmoduleImport(SubmoduleImport { call_path }) => { - Some(Cow::Owned(format_call_path(call_path))) - } - _ => 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) @@ -249,51 +205,6 @@ impl<'a> Binding<'a> { } } -pub enum AnyImport<'a> { - Import(&'a Import<'a>), - SubmoduleImport(&'a SubmoduleImport<'a>), - FromImport(&'a FromImport<'a>), -} - -impl<'a> AnyImport<'a> { - /// Returns the fully-qualified symbol name, if this symbol was imported from another module. - pub fn call_path(&self) -> &[&str] { - match self { - Self::Import(Import { call_path }) => call_path.as_ref(), - Self::SubmoduleImport(SubmoduleImport { call_path }) => call_path.as_ref(), - Self::FromImport(FromImport { call_path }) => call_path.as_ref(), - } - } - - /// Returns the fully-qualified symbol name, if this symbol was imported from another module. - pub fn qualified_name(&self) -> String { - format_call_path(self.call_path()) - } - - /// 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) -> &[&str] { - match self { - Self::Import(Import { call_path }) => &call_path[..1], - Self::SubmoduleImport(SubmoduleImport { call_path }) => &call_path[..1], - Self::FromImport(FromImport { call_path }) => &call_path[..call_path.len() - 1], - } - } - - /// Returns the fully-qualified symbol name, if this symbol was imported from another module. - pub fn member_name(&self) -> Cow<'a, str> { - match self { - Self::Import(Import { call_path }) => Cow::Owned(format_call_path(call_path)), - Self::SubmoduleImport(SubmoduleImport { call_path }) => { - Cow::Owned(format_call_path(call_path)) - } - Self::FromImport(FromImport { call_path }) => { - Cow::Borrowed(call_path[call_path.len() - 1]) - } - } - } -} - bitflags! { /// Flags on a [`Binding`]. #[derive(Debug, Default, Copy, Clone, Eq, PartialEq)] @@ -603,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 [`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 39685c5e56..77c44f0759 100644 --- a/crates/ruff_python_semantic/src/model.rs +++ b/crates/ruff_python_semantic/src/model.rs @@ -4,9 +4,8 @@ use bitflags::bitflags; use rustc_hash::FxHashMap; use smallvec::smallvec; -use ruff_python_ast::call_path::{ - collect_call_path, collect_head_path, from_unqualified_name, CallPath, -}; +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; @@ -24,11 +23,12 @@ 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> { typing_modules: &'a [String], + module_path: Option<&'a [String]>, /// Stack of all visited statements. pub stmts: Nodes<'a>, @@ -128,6 +128,7 @@ impl<'a> SemanticModel<'a> { pub fn new(typing_modules: &'a [String], path: &'a Path, module: Module<'a>) -> Self { Self { typing_modules, + module_path: module.path(), stmts: Nodes::default(), stmt_id: None, exprs: Vec::default(), @@ -583,7 +584,8 @@ impl<'a> SemanticModel<'a> { // import pyarrow.csv // print(pa.csv.read_csv("test.csv")) // ``` - let call_path = self.bindings[binding_id].call_path()?; + 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; @@ -619,9 +621,8 @@ impl<'a> SemanticModel<'a> { } } - let (head, tail) = collect_head_path(value)?; - // If the name was already resolved, look it up; otherwise, search for the symbol. + let head = match_head(value)?; let binding = if let Some(id) = self.resolved_names.get(&head.into()) { self.binding(*id) } else { @@ -630,16 +631,32 @@ impl<'a> SemanticModel<'a> { match &binding.kind { 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 { call_path }) => { - let mut source_path: CallPath = CallPath::from_slice(&call_path[..1]); - source_path.extend_from_slice(tail.as_slice()); - Some(source_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 { call_path }) => { - let resolved: CallPath = call_path.iter().chain(tail.iter()).copied().collect(); + let value_path = collect_call_path(value)?; + let (_, tail) = value_path.split_first()?; + + let resolved: CallPath = + if call_path.first().map_or(false, |segment| *segment == ".") { + from_relative_import(self.module_path?, call_path, tail)? + } else { + call_path.iter().chain(tail.iter()).copied().collect() + }; Some(resolved) } BindingKind::Builtin => Some(smallvec!["", head.id.as_str()]),