Merge pull request #42 from pylint-dev/21-add-pandas-convert-df-to-numpy-checker

Add pandas dataframe values checker

Author: PeterHamfelt

pylint_ml/checkers/pandas/pandas_dataframe_values.py

@@ -0,0 +1,30 @@
+# Licensed under the MIT: https://mit-license.org/
+# For details: https://github.com/pylint-dev/pylint-ml/LICENSE
+# Copyright (c) https://github.com/pylint-dev/pylint-ml/CONTRIBUTORS.txt
+
+"""Check for the usage of pandas.DataFrame.values and suggest .to_numpy() instead."""
+
+from __future__ import annotations
+
+from astroid import nodes
+from pylint.checkers import BaseChecker
+from pylint.checkers.utils import only_required_for_messages
+from pylint.interfaces import HIGH
+
+
+class PandasValuesChecker(BaseChecker):
+    name = "pandas-dataframe-values"
+    msgs = {
+        "W8112": (
+            "Avoid using 'DataFrame.values'. Use '.to_numpy()' instead for better consistency and compatibility.",
+            "pandas-dataframe-values",
+            "Using 'DataFrame.values' is discouraged as it may not always return a NumPy array. Use '.to_numpy()' "
+            "instead.",
+        ),
+    }
+
+    @only_required_for_messages("pandas-dataframe-values")
+    def visit_attribute(self, node: nodes.Attribute) -> None:
+        if isinstance(node.expr, nodes.Name):
+            if node.attrname == "values" and node.expr.name.startswith("df_"):
+                self.add_message("pandas-dataframe-values", node=node, confidence=HIGH)

tests/checkers/test_pandas/test_pandas_dataframe_values.py

@@ -0,0 +1,49 @@
+import astroid
+import pylint.testutils
+from pylint.interfaces import HIGH
+
+from pylint_ml.checkers.pandas.pandas_dataframe_values import PandasValuesChecker
+
+
+class TestPandasValuesChecker(pylint.testutils.CheckerTestCase):
+    CHECKER_CLASS = PandasValuesChecker
+
+    def test_values_usage_with_correct_naming(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            df_sales = pd.DataFrame({
+                "A": [1, 2, 3],
+                "B": [4, 5, 6]
+            })
+            data = df_sales.values  # [pandas-dataframe-values]
+            """
+        )
+
+        # Access the attribute that is 'values'
+        attribute_node = node.value
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="pandas-dataframe-values",
+                confidence=HIGH,
+                node=attribute_node,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_attribute(attribute_node)
+
+    def test_no_warning_for_to_numpy(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            df_sales = pd.DataFrame({
+                "A": [1, 2, 3],
+                "B": [4, 5, 6]
+            })
+            df_data = df_sales.to_numpy()  # This should not trigger any warnings
+            """
+        )
+
+        with self.assertNoMessages():
+            self.checker.visit_call(node)