diff --git a/crates/ruff_linter/src/rules/pandas_vet/mod.rs b/crates/ruff_linter/src/rules/pandas_vet/mod.rs index 982eedbde5..3e96328c69 100644 --- a/crates/ruff_linter/src/rules/pandas_vet/mod.rs +++ b/crates/ruff_linter/src/rules/pandas_vet/mod.rs @@ -269,6 +269,45 @@ mod tests { ", "PD011_pass_node_name" )] + #[test_case( + r" + import pandas as pd + import numpy as np + unique = np.unique_inverse([1, 2, 3, 2, 1]) + result = unique.values + ", + "PD011_pass_numpy_unique_inverse" + )] + #[test_case( + r" + import pandas as pd + import numpy as np + unique = np.unique_all([1, 2, 3, 2, 1]) + result = unique.values + ", + "PD011_pass_numpy_unique_all" + )] + #[test_case( + r" + import pandas as pd + import numpy as np + unique = np.unique_counts([1, 2, 3, 2, 1]) + result = unique.values + ", + "PD011_pass_numpy_unique_counts" + )] + #[test_case( + r" + import pandas as pd + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from numpy.lib._arraysetops_impl import UniqueInverseResult + import numpy as np + unique: UniqueInverseResult[np.uint64] = np.unique_inverse([1, 2, 3, 2, 1]) + result = unique.values + ", + "PD011_pass_numpy_typed_unique_inverse" + )] #[test_case( r#" import pandas as pd diff --git a/crates/ruff_linter/src/rules/pandas_vet/rules/attr.rs b/crates/ruff_linter/src/rules/pandas_vet/rules/attr.rs index ab6b55104b..c5f34c94ed 100644 --- a/crates/ruff_linter/src/rules/pandas_vet/rules/attr.rs +++ b/crates/ruff_linter/src/rules/pandas_vet/rules/attr.rs @@ -1,6 +1,6 @@ use ruff_macros::{ViolationMetadata, derive_message_formats}; use ruff_python_ast::{self as ast, Expr}; -use ruff_python_semantic::Modules; +use ruff_python_semantic::{Modules, analyze::typing::find_binding_value}; use ruff_text_size::Ranged; use crate::Violation; @@ -43,6 +43,38 @@ impl Violation for PandasUseOfDotValues { } } +/// Check if a binding comes from a NumPy function that returns a `NamedTuple` with a `.values` field. +fn is_numpy_namedtuple_binding( + expr: &Expr, + semantic: &ruff_python_semantic::SemanticModel, +) -> bool { + let Expr::Name(name) = expr else { + return false; + }; + + let Some(binding_id) = semantic.resolve_name(name) else { + return false; + }; + let binding = semantic.binding(binding_id); + + let Some(assigned_value) = find_binding_value(binding, semantic) else { + return false; + }; + + let Some(call_expr) = assigned_value.as_call_expr() else { + return false; + }; + + let Some(qualified_name) = semantic.resolve_qualified_name(&call_expr.func) else { + return false; + }; + + matches!( + qualified_name.segments(), + ["numpy", "unique_inverse" | "unique_all" | "unique_counts"] + ) +} + /// PD011 pub(crate) fn attr(checker: &Checker, attribute: &ast::ExprAttribute) { if !checker.semantic().seen_module(Modules::PANDAS) { @@ -77,5 +109,10 @@ pub(crate) fn attr(checker: &Checker, attribute: &ast::ExprAttribute) { return; } + // Avoid flagging on NumPy `NamedTuples` that have a legitimate `.values` field + if is_numpy_namedtuple_binding(attribute.value.as_ref(), checker.semantic()) { + return; + } + checker.report_diagnostic(PandasUseOfDotValues, attribute.range()); } diff --git a/crates/ruff_linter/src/rules/pandas_vet/snapshots/ruff_linter__rules__pandas_vet__tests__PD011_pass_numpy_typed_unique_inverse.snap b/crates/ruff_linter/src/rules/pandas_vet/snapshots/ruff_linter__rules__pandas_vet__tests__PD011_pass_numpy_typed_unique_inverse.snap new file mode 100644 index 0000000000..07173f8f1b --- /dev/null +++ b/crates/ruff_linter/src/rules/pandas_vet/snapshots/ruff_linter__rules__pandas_vet__tests__PD011_pass_numpy_typed_unique_inverse.snap @@ -0,0 +1,4 @@ +--- +source: crates/ruff_linter/src/rules/pandas_vet/mod.rs +--- + diff --git a/crates/ruff_linter/src/rules/pandas_vet/snapshots/ruff_linter__rules__pandas_vet__tests__PD011_pass_numpy_unique_all.snap b/crates/ruff_linter/src/rules/pandas_vet/snapshots/ruff_linter__rules__pandas_vet__tests__PD011_pass_numpy_unique_all.snap new file mode 100644 index 0000000000..07173f8f1b --- /dev/null +++ b/crates/ruff_linter/src/rules/pandas_vet/snapshots/ruff_linter__rules__pandas_vet__tests__PD011_pass_numpy_unique_all.snap @@ -0,0 +1,4 @@ +--- +source: crates/ruff_linter/src/rules/pandas_vet/mod.rs +--- + diff --git a/crates/ruff_linter/src/rules/pandas_vet/snapshots/ruff_linter__rules__pandas_vet__tests__PD011_pass_numpy_unique_counts.snap b/crates/ruff_linter/src/rules/pandas_vet/snapshots/ruff_linter__rules__pandas_vet__tests__PD011_pass_numpy_unique_counts.snap new file mode 100644 index 0000000000..07173f8f1b --- /dev/null +++ b/crates/ruff_linter/src/rules/pandas_vet/snapshots/ruff_linter__rules__pandas_vet__tests__PD011_pass_numpy_unique_counts.snap @@ -0,0 +1,4 @@ +--- +source: crates/ruff_linter/src/rules/pandas_vet/mod.rs +--- + diff --git a/crates/ruff_linter/src/rules/pandas_vet/snapshots/ruff_linter__rules__pandas_vet__tests__PD011_pass_numpy_unique_inverse.snap b/crates/ruff_linter/src/rules/pandas_vet/snapshots/ruff_linter__rules__pandas_vet__tests__PD011_pass_numpy_unique_inverse.snap new file mode 100644 index 0000000000..07173f8f1b --- /dev/null +++ b/crates/ruff_linter/src/rules/pandas_vet/snapshots/ruff_linter__rules__pandas_vet__tests__PD011_pass_numpy_unique_inverse.snap @@ -0,0 +1,4 @@ +--- +source: crates/ruff_linter/src/rules/pandas_vet/mod.rs +--- +