Add dataframe and series checkers

Author: Peter Hamfelt

pylint_ml/checkers/pandas/pandas_dataframe_bool.py

@@ -11,6 +11,8 @@
 from pylint.checkers.utils import only_required_for_messages
 from pylint.interfaces import HIGH
 
+# Todo add version deprecated
+
 
 class PandasDataFrameBoolChecker(BaseChecker):
     name = "pandas-dataframe-bool"
@@ -26,13 +28,14 @@ class PandasDataFrameBoolChecker(BaseChecker):
     def visit_call(self, node: nodes.Call) -> None:
         if isinstance(node.func, nodes.Attribute):
             method_name = getattr(node.func, "attrname", None)
-            module_name = getattr(node.func.expr, "name", None)
 
-            if method_name == "bool" and module_name == "pd":
-                self.add_message("pandas-dataframe-bool", node=node, confidence=HIGH)
+            if method_name == "bool":
+                # Check if the object calling .bool() has a name starting with 'df_'
+                object_name = getattr(node.func.expr, "name", None)
+                if object_name and self._is_valid_dataframe_name(object_name):
+                    self.add_message("pandas-dataframe-bool", node=node, confidence=HIGH)
 
-    def _check_method_usage(self, node):
-        method_name = getattr(node.func, "attrname", None)
-        module_name = getattr(node.func.expr, "name", None)
-        if method_name == "bool" and module_name == "pd":
-            self.add_message("pandas-dataframe-bool", node=node, confidence=HIGH)
+    @staticmethod
+    def _is_valid_dataframe_name(name: str) -> bool:
+        """Check if the DataFrame name starts with 'df_'."""
+        return name.startswith("df_")

pylint_ml/checkers/pandas/pandas_dataframe_naming.py

@@ -29,13 +29,8 @@ def visit_assign(self, node: nodes.Assign) -> None:
             module_name = getattr(node.value.func.expr, "name", None)
 
             if func_name == "DataFrame" and module_name == "pd":
-
                 for target in node.targets:
                     if isinstance(target, nodes.AssignName):
                         var_name = target.name
                         if not var_name.startswith("df_") or len(var_name) <= 3:
                             self.add_message("pandas-dataframe-naming", node=node, confidence=HIGH)
-
-    def _check_variable_name(self, var_name, node):
-        if not var_name.startswith("df_") or len(var_name) <= 3:
-            self.add_message("pandas-dataframe-naming", node=node, confidence=HIGH)

pylint_ml/checkers/pandas/pandas_series_bool.py

@@ -11,6 +11,8 @@
 from pylint.checkers.utils import only_required_for_messages
 from pylint.interfaces import HIGH
 
+# Todo add version deprecated
+
 
 class PandasSeriesBoolChecker(BaseChecker):
     name = "pandas-series-bool"
@@ -26,13 +28,14 @@ class PandasSeriesBoolChecker(BaseChecker):
     def visit_call(self, node: nodes.Call) -> None:
         if isinstance(node.func, nodes.Attribute):
             method_name = getattr(node.func, "attrname", None)
-            module_name = getattr(node.func.expr, "name", None)
 
-            if method_name == "bool" and module_name == "pd":
-                self.add_message("pandas-series-bool", node=node, confidence=HIGH)
+            if method_name == "bool":
+                # Check if the object calling .bool() has a name starting with 'ser'
+                object_name = getattr(node.func.expr, "name", None)
+                if object_name and self._is_valid_series_name(object_name):
+                    self.add_message("pandas-series-bool", node=node, confidence=HIGH)
 
-    def _check_method_usage(self, node):
-        method_name = getattr(node.func, "attrname", None)
-        module_name = getattr(node.func.expr, "name", None)
-        if method_name == "bool" and module_name == "pd":
-            self.add_message("pandas-series-bool", node=node, confidence=HIGH)
+    @staticmethod
+    def _is_valid_series_name(name: str) -> bool:
+        """Check if the Series name starts with 'ser_'."""
+        return name.startswith("ser_")

pylint_ml/checkers/pandas/pandas_series_naming.py

@@ -0,0 +1,37 @@
+# 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 consistent naming of pandas Series variables."""
+
+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 PandasSeriesNamingChecker(BaseChecker):
+    name = "pandas-series-naming"
+    msgs = {
+        "W8103": (
+            "Pandas Series variable names should start with 'ser_' followed by descriptive text",
+            "pandas-series-naming",
+            "Ensure that pandas Series variables follow the naming convention.",
+        ),
+    }
+
+    @only_required_for_messages("pandas-series-naming")
+    def visit_assign(self, node: nodes.Assign) -> None:
+        print(node)
+        if isinstance(node.value, nodes.Call):
+            func_name = getattr(node.value.func, "attrname", None)
+            module_name = getattr(node.value.func.expr, "name", None)
+
+            if func_name == "Series" and module_name == "pd":
+                for target in node.targets:
+                    if isinstance(target, nodes.AssignName):
+                        var_name = target.name
+                        if not var_name.startswith("ser_") or len(var_name) <= 4:
+                            self.add_message("pandas-series-naming", node=node, confidence=HIGH)

tests/checkers/test_pandas/test_pandas_dataframe_bool.py

@@ -21,7 +21,8 @@ def test_dataframe_bool_usage(self):
                 msg_id="pandas-dataframe-bool",
                 confidence=HIGH,
                 node=node,
-            )
+            ),
+            ignore_position=True,
         ):
             self.checker.visit_call(node)
 

tests/checkers/test_pandas/test_dataframe_naming.py renamed to tests/checkers/test_pandas/test_pandas_dataframe_naming.py

@@ -0,0 +1,53 @@
+import astroid
+import pylint.testutils
+from pylint.interfaces import HIGH
+
+from pylint_ml.checkers.pandas.pandas_series_naming import PandasSeriesNamingChecker
+
+
+class TestPandasSeriesNamingChecker(pylint.testutils.CheckerTestCase):
+    CHECKER_CLASS = PandasSeriesNamingChecker
+
+    def test_series_correct_naming(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            ser_sales = pd.Series([100, 200, 300])
+            """
+        )
+        with self.assertNoMessages():
+            self.checker.visit_assign(node)
+
+    def test_series_incorrect_naming(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            df_sales = pd.Series([100, 200, 300])
+            """
+        )
+        with self.assertAddsMessages(
+                pylint.testutils.MessageTest(
+                    msg_id="pandas-series-naming",
+                    confidence=HIGH,
+                    node=node,
+                ),
+                ignore_position=True,
+        ):
+            self.checker.visit_assign(node)
+
+    def test_series_invalid_length_naming(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            ser_ = pd.Series([True])
+            """
+        )
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="pandas-series-naming",
+                confidence=HIGH,
+                node=node,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_assign(node)