Merge pull request #41 from pylint-dev/25-add-pandas-dataframe-merge-parameter-not-explicitly-set-checker

Add pandas dataframe merge parameter checker

Author: PeterHamfelt

pylint_ml/checkers/pandas/pandas_dataframe_merge.py

@@ -0,0 +1,41 @@
+# 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 proper usage of pandas merge() with explicit parameters and DataFrame naming conventions."""
+
+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 PandasDataframeMergeChecker(BaseChecker):
+    name = "pandas-dataframe-merge"
+    msgs = {
+        "W8110": (
+            "Ensure that 'how', 'on', and 'validate' parameters are explicitly specified in pandas DataFrame merge().",
+            "pandas-dataframe-merge",
+            "Explicitly specifying the 'how', 'on', and 'validate' parameters and using a proper DataFrame naming "
+            "convention improves readability and prevents unintended behavior.",
+        ),
+    }
+
+    @only_required_for_messages("pandas-dataframe-merge")
+    def visit_call(self, node: nodes.Call) -> None:
+        if isinstance(node.func, nodes.Attribute):
+            method_name = node.func.attrname
+            object_name = getattr(node.func.expr, "name", "")
+            if method_name == "merge":
+                # Check if the DataFrame name starts with 'df_'
+                name_is_valid = object_name.startswith("df_")
+
+                # Check for explicit 'how', 'on', and 'validate' parameters
+                how_specified = any(kw.arg == "how" for kw in node.keywords)
+                on_specified = any(kw.arg == "on" for kw in node.keywords)
+                validate_specified = any(kw.arg == "validate" for kw in node.keywords)
+
+                if not (name_is_valid and how_specified and on_specified and validate_specified):
+                    self.add_message("pandas-dataframe-merge", node=node, confidence=HIGH)

tests/checkers/test_pandas/test_pandas_dataframe_merge.py

@@ -0,0 +1,82 @@
+import astroid
+import pylint.testutils
+from pylint.interfaces import HIGH
+
+from pylint_ml.checkers.pandas.pandas_dataframe_merge import PandasDataframeMergeChecker
+
+
+class TestPandasMergeChecker(pylint.testutils.CheckerTestCase):
+    CHECKER_CLASS = PandasDataframeMergeChecker
+
+    def test_merge_without_explicit_params(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            df_3 = df_1.merge(df_2)  # [pandas-dataframe-merge]
+            """
+        )
+
+        merge_call = node.value
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="pandas-dataframe-merge",
+                confidence=HIGH,
+                node=merge_call,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_call(merge_call)
+
+    def test_merge_with_missing_validate(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            df_3 = df_1.merge(df_2, how='inner', on='col1')  # [pandas-dataframe-merge]
+            """
+        )
+
+        merge_call = node.value
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="pandas-dataframe-merge",
+                confidence=HIGH,
+                node=merge_call,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_call(merge_call)
+
+    def test_merge_with_wrong_naming_and_missing_params(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            merged_df = df_1.merge(df_2)  # [pandas-dataframe-merge]
+            """
+        )
+
+        merge_call = node.value
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="pandas-dataframe-merge",
+                confidence=HIGH,
+                node=merge_call,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_call(merge_call)
+
+    def test_merge_with_all_params_and_correct_naming(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            df_merged = df_1.merge(df_2, how='inner', on='col1', validate='1:1')  # This should not trigger any warnings
+            """
+        )
+
+        merge_call = node.value
+
+        with self.assertNoMessages():
+            self.checker.visit_call(merge_call)