Merge pull request #43 from pylint-dev/20-add-pandas-empty-column-misinitialization-checker

Add pandas dataframe empty column checker

Author: PeterHamfelt

pylint_ml/checkers/pandas/pandas_dataframe_empty_column.py

@@ -0,0 +1,33 @@
+# 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 initialization of empty columns in pandas DataFrames."""
+
+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 PandasEmptyColumnChecker(BaseChecker):
+    name = "pandas-dataframe-empty-column"
+    msgs = {
+        "W8113": (
+            "Avoid using filler values (0, '') for new empty columns. Use 'np.nan' or 'pd.Series(dtype=...)' instead.",
+            "pandas-dataframe-empty-column",
+            "Initializing new columns with filler values such as 0 or empty strings can lead to issues with null "
+            "value detection.",
+        ),
+    }
+
+    @only_required_for_messages("pandas-dataframe-empty-column")
+    def visit_subscript(self, node: nodes.Subscript) -> None:
+        if isinstance(node.value, nodes.Name) and node.value.name.startswith("df_"):
+            if isinstance(node.slice, nodes.Const) and isinstance(node.parent, nodes.Assign):
+                if isinstance(node.parent.value, nodes.Const):
+                    # Checking for filler values: 0 or empty string
+                    if node.parent.value.value in (0, ""):
+                        self.add_message("pandas-dataframe-empty-column", node=node, confidence=HIGH)

tests/checkers/test_pandas/test_pandas_dataframe_empty_column.py

@@ -0,0 +1,68 @@
+import astroid
+import pylint.testutils
+from pylint.interfaces import HIGH
+
+from pylint_ml.checkers.pandas.pandas_dataframe_empty_column import PandasEmptyColumnChecker
+
+
+class TestPandasEmptyColumnChecker(pylint.testutils.CheckerTestCase):
+    CHECKER_CLASS = PandasEmptyColumnChecker
+
+    def test_correct_empty_column_initialization(self):
+        node = astroid.extract_node(
+            """
+            import numpy as np
+            import pandas as pd
+            df_sales = pd.DataFrame()
+            df_sales['new_col_float'] = np.nan  # This should not trigger any warnings
+            df_sales['new_col_int'] = pd.Series(dtype='int')  # This should not trigger any warnings
+            df_sales['new_col_str'] = pd.Series(dtype='object')  # This should not trigger any warnings
+            """
+        )
+        with self.assertNoMessages():
+            self.checker.visit_subscript(node)
+            self.checker.visit_subscript(node)
+            self.checker.visit_subscript(node)
+
+    def test_incorrect_empty_column_initialization_with_zero(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            df_sales = pd.DataFrame()
+            df_sales['new_col_int'] = 0  # [pandas-dataframe-empty-column]
+            """
+        )
+
+        # Extract the Subscript node for the df_sales['new_col_int'] assignment
+        subscript_node = node.targets[0]
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="pandas-dataframe-empty-column",
+                confidence=HIGH,
+                node=subscript_node,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_subscript(subscript_node)
+
+    def test_incorrect_empty_column_initialization_with_empty_string(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            df_sales = pd.DataFrame()
+            df_sales['new_col_str'] = ''  # [pandas-dataframe-empty-column]
+            """
+        )
+
+        subscript_node = node.targets[0]
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="pandas-dataframe-empty-column",
+                confidence=HIGH,
+                node=subscript_node,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_subscript(subscript_node)