diff --git a/src/checkers/ast.rs b/src/checkers/ast.rs index a6044e265f..e4ec6608f7 100644 --- a/src/checkers/ast.rs +++ b/src/checkers/ast.rs @@ -1977,49 +1977,7 @@ where flake8_2020::rules::name_or_attribute(self, expr); } - for (code, name) in vec![ - (Rule::UseOfDotIx, "ix"), - (Rule::UseOfDotAt, "at"), - (Rule::UseOfDotIat, "iat"), - (Rule::UseOfDotValues, "values"), - ] { - if self.settings.rules.enabled(&code) { - if attr == name { - // Avoid flagging on function calls (e.g., `df.values()`). - if let Some(parent) = self.current_expr_parent() { - if matches!(parent.node, ExprKind::Call { .. }) { - continue; - } - } - // Avoid flagging on non-DataFrames (e.g., `{"a": 1}.values`). - if pandas_vet::helpers::is_dataframe_candidate(value) { - // If the target is a named variable, avoid triggering on - // irrelevant bindings (like imports). - if let ExprKind::Name { id, .. } = &value.node { - if self.find_binding(id).map_or(true, |binding| { - matches!( - binding.kind, - BindingKind::Builtin - | BindingKind::ClassDefinition - | BindingKind::FunctionDefinition - | BindingKind::Export(..) - | BindingKind::FutureImportation - | BindingKind::StarImportation(..) - | BindingKind::Importation(..) - | BindingKind::FromImportation(..) - | BindingKind::SubmoduleImportation(..) - ) - }) { - continue; - } - } - - self.diagnostics - .push(Diagnostic::new(code.kind(), Range::from_located(expr))); - } - }; - } - } + pandas_vet::rules::check_attr(self, attr, value, expr); if self.settings.rules.enabled(&Rule::BannedApi) { flake8_tidy_imports::banned_api::banned_attribute_access(self, expr); @@ -2377,54 +2335,8 @@ where self.diagnostics .extend(pandas_vet::rules::inplace_argument(keywords).into_iter()); } - for (code, name) in vec![ - (Rule::UseOfDotIsNull, "isnull"), - (Rule::UseOfDotNotNull, "notnull"), - (Rule::UseOfDotPivotOrUnstack, "pivot"), - (Rule::UseOfDotPivotOrUnstack, "unstack"), - (Rule::UseOfDotReadTable, "read_table"), - (Rule::UseOfDotStack, "stack"), - ] { - if self.settings.rules.enabled(&code) { - if let ExprKind::Attribute { value, attr, .. } = &func.node { - if attr == name { - if pandas_vet::helpers::is_dataframe_candidate(value) { - // If the target is a named variable, avoid triggering on - // irrelevant bindings (like non-Pandas imports). - if let ExprKind::Name { id, .. } = &value.node { - if self.find_binding(id).map_or(true, |binding| { - if let BindingKind::Importation(.., module) = - &binding.kind - { - module != &"pandas" - } else { - matches!( - binding.kind, - BindingKind::Builtin - | BindingKind::ClassDefinition - | BindingKind::FunctionDefinition - | BindingKind::Export(..) - | BindingKind::FutureImportation - | BindingKind::StarImportation(..) - | BindingKind::Importation(..) - | BindingKind::FromImportation(..) - | BindingKind::SubmoduleImportation(..) - ) - } - }) { - continue; - } - } + pandas_vet::rules::check_call(self, func); - self.diagnostics.push(Diagnostic::new( - code.kind(), - Range::from_located(func), - )); - } - }; - } - } - } if self.settings.rules.enabled(&Rule::UseOfPdMerge) { if let Some(diagnostic) = pandas_vet::rules::use_of_pd_merge(func) { self.diagnostics.push(diagnostic); diff --git a/src/rules/pandas_vet/rules.rs b/src/rules/pandas_vet/rules.rs index 066d34fe30..8244bf676f 100644 --- a/src/rules/pandas_vet/rules.rs +++ b/src/rules/pandas_vet/rules.rs @@ -1,7 +1,8 @@ -use rustpython_ast::{Constant, Expr, ExprKind, Keyword}; +use rustpython_ast::{Constant, Expr, ExprKind, Keyword, Located}; -use crate::ast::types::Range; -use crate::registry::Diagnostic; +use crate::ast::types::{BindingKind, Range}; +use crate::checkers::ast::Checker; +use crate::registry::{Diagnostic, Rule}; use crate::violations; /// PD002 @@ -60,3 +61,103 @@ pub fn assignment_to_df(targets: &[Expr]) -> Option { Range::from_located(target), )) } + +pub fn check_attr( + checker: &mut Checker, + attr: &str, + value: &Located, + attr_expr: &Located, +) { + for (code, name) in vec![ + (Rule::UseOfDotIx, "ix"), + (Rule::UseOfDotAt, "at"), + (Rule::UseOfDotIat, "iat"), + (Rule::UseOfDotValues, "values"), + ] { + if checker.settings.rules.enabled(&code) { + if attr == name { + // Avoid flagging on function calls (e.g., `df.values()`). + if let Some(parent) = checker.current_expr_parent() { + if matches!(parent.node, ExprKind::Call { .. }) { + continue; + } + } + // Avoid flagging on non-DataFrames (e.g., `{"a": 1}.values`). + if super::helpers::is_dataframe_candidate(value) { + // If the target is a named variable, avoid triggering on + // irrelevant bindings (like imports). + if let ExprKind::Name { id, .. } = &value.node { + if checker.find_binding(id).map_or(true, |binding| { + matches!( + binding.kind, + BindingKind::Builtin + | BindingKind::ClassDefinition + | BindingKind::FunctionDefinition + | BindingKind::Export(..) + | BindingKind::FutureImportation + | BindingKind::StarImportation(..) + | BindingKind::Importation(..) + | BindingKind::FromImportation(..) + | BindingKind::SubmoduleImportation(..) + ) + }) { + continue; + } + } + + checker + .diagnostics + .push(Diagnostic::new(code.kind(), Range::from_located(attr_expr))); + } + }; + } + } +} + +pub fn check_call(checker: &mut Checker, func: &Located) { + for (code, name) in vec![ + (Rule::UseOfDotIsNull, "isnull"), + (Rule::UseOfDotNotNull, "notnull"), + (Rule::UseOfDotPivotOrUnstack, "pivot"), + (Rule::UseOfDotPivotOrUnstack, "unstack"), + (Rule::UseOfDotReadTable, "read_table"), + (Rule::UseOfDotStack, "stack"), + ] { + if checker.settings.rules.enabled(&code) { + if let ExprKind::Attribute { value, attr, .. } = &func.node { + if attr == name { + if super::helpers::is_dataframe_candidate(value) { + // If the target is a named variable, avoid triggering on + // irrelevant bindings (like non-Pandas imports). + if let ExprKind::Name { id, .. } = &value.node { + if checker.find_binding(id).map_or(true, |binding| { + if let BindingKind::Importation(.., module) = &binding.kind { + module != &"pandas" + } else { + matches!( + binding.kind, + BindingKind::Builtin + | BindingKind::ClassDefinition + | BindingKind::FunctionDefinition + | BindingKind::Export(..) + | BindingKind::FutureImportation + | BindingKind::StarImportation(..) + | BindingKind::Importation(..) + | BindingKind::FromImportation(..) + | BindingKind::SubmoduleImportation(..) + ) + } + }) { + continue; + } + } + + checker + .diagnostics + .push(Diagnostic::new(code.kind(), Range::from_located(func))); + } + }; + } + } + } +}