refactor: Move a bunch of pandas-vet logic to rules::pandas_vet

This commit is contained in:
Martin Fischer 2023-01-19 12:00:06 +01:00 committed by Charlie Marsh
parent a122d95ef5
commit 9603a024b3
2 changed files with 106 additions and 93 deletions

View File

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

View File

@ -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<Diagnostic> {
Range::from_located(target),
))
}
pub fn check_attr(
checker: &mut Checker,
attr: &str,
value: &Located<ExprKind>,
attr_expr: &Located<ExprKind>,
) {
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<ExprKind>) {
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)));
}
};
}
}
}
}