Merge pull request #44 from pylint-dev/19-add-pandas-column-datatype-not-explicitly-set-checker

Add pandas dtype parameter checker

Author: PeterHamfelt

pylint_ml/checkers/pandas/pandas_copy_on_write.py

@@ -0,0 +1,27 @@
+from astroid import nodes
+from pylint.checkers import BaseChecker
+from pylint.checkers.utils import only_required_for_messages
+from pylint.interfaces import HIGH
+
+
+class PandasDtypeChecker(BaseChecker):
+    name = "pandas-dtype-param"
+    msgs = {
+        "W8117": (
+            "Specify 'dtype' when using '%s' for better performance and data integrity.",
+            "pandas-dtype-param",
+            "It's recommended to explicitly specify the 'dtype' parameter in pandas read functions.",
+        ),
+    }
+
+    @only_required_for_messages("pandas-dtype-param")
+    def visit_call(self, node: nodes.Call) -> None:
+        # Check if the function being called is a pandas read function
+        if isinstance(node.func, nodes.Attribute):
+            module_name = getattr(node.func.expr, "name", None)
+            func_name = node.func.attrname
+
+            if module_name == "pd" and func_name in {"read_csv", "read_excel", "read_table"}:
+                # Check if dtype is specified
+                if not any(kw.arg == "dtype" for kw in node.keywords):
+                    self.add_message("pandas-dtype-param", node=node, confidence=HIGH)

pylint_ml/checkers/pandas/pandas_dtype_param.py

@@ -32,18 +32,3 @@ def test_values_usage_with_correct_naming(self):
             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)

tests/checkers/test_pandas/test_pandas_copy_on_write.py

@@ -0,0 +1,42 @@
+import astroid
+import pylint.testutils
+from pylint.interfaces import HIGH
+
+from pylint_ml.checkers.pandas.pandas_dtype_param import PandasDtypeChecker
+
+
+class TestPandasDtypeChecker(pylint.testutils.CheckerTestCase):
+    CHECKER_CLASS = PandasDtypeChecker
+
+    def test_dtype_specified(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            df = pd.read_csv('file.csv', dtype={'column1': 'int32'})
+            """
+        )
+
+        dtype_call = node.value
+
+        with self.assertNoMessages():
+            self.checker.visit_call(dtype_call)
+
+    def test_dtype_missing(self):
+        node = astroid.extract_node(
+            """
+            import pandas as pd
+            df_sales = pd.read_csv('file.csv')
+            """
+        )
+
+        dtype_call = node.value
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="pandas-dtype-param",
+                confidence=HIGH,
+                node=dtype_call,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_call(dtype_call)