Merge pull request #38 from pylint-dev/17-add-unnecessary-iteration-avoidance-pandas-checker

Add pandas iterrows checker

Author: PeterHamfelt

pylint_ml/checkers/pandas/pandas_dataframe_iterrows.py

@@ -0,0 +1,38 @@
+# 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 usage of the inefficient pandas DataFrame.iterrows() method."""
+
+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 PandasIterrowsChecker(BaseChecker):
+    name = "pandas-iterrows"
+    msgs = {
+        "W8106": (
+            "Usage of pandas DataFrame.iterrows() detected",
+            "pandas-iterrows",
+            "Avoid using DataFrame.iterrows() for large datasets. Consider using vectorized operations or "
+            ".itertuples() instead.",
+        ),
+    }
+
+    @only_required_for_messages("pandas-iterrows")
+    def visit_call(self, node: nodes.Call) -> None:
+        if isinstance(node.func, nodes.Attribute):
+            method_name = getattr(node.func, "attrname", None)
+            if method_name == "iterrows":
+                object_name = getattr(node.func.expr, "name", None)
+                if object_name and self._is_dataframe_name(object_name):
+                    self.add_message("pandas-iterrows", node=node, confidence=HIGH)
+
+    @staticmethod
+    def _is_dataframe_name(name: str) -> bool:
+        """Check if the object name suggests it's a DataFrame (e.g., starts with 'df_')."""
+        return name.startswith("df_")

tests/checkers/test_pandas/test_pandas_dataframe_iterrows.py

@@ -0,0 +1,35 @@
+import astroid
+import pylint.testutils
+from pylint.interfaces import HIGH
+
+from pylint_ml.checkers.pandas.pandas_dataframe_iterrows import PandasIterrowsChecker
+
+
+class TestPandasIterrowsChecker(pylint.testutils.CheckerTestCase):
+    CHECKER_CLASS = PandasIterrowsChecker
+
+    def test_iterrows_used(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            df_sales = pd.DataFrame({
+                "Product": ["A", "B", "C"],
+                "Sales": [100, 200, 300]
+            })
+            for index, row in df_sales.iterrows():  # [pandas-iterrows]
+                print(row["Product"], row["Sales"])
+            """
+        )
+
+        # Extract the Call node for the `iterrows` method
+        iterrows_call = node.iter  # This directly points to the `Call` node for `iterrows()`
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="pandas-iterrows",
+                confidence=HIGH,
+                node=iterrows_call,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_call(iterrows_call)