Add numpy nan comparison checker

Author: PeterHamfelt

pylint_ml/checkers/numpy/import_numpy.py renamed to pylint_ml/checkers/numpy/numpy_import.py

@@ -0,0 +1,50 @@
+# 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 numpy nan comparison."""
+
+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
+
+COMPARISON_OP = frozenset(("<", "<=", ">", ">=", "!=", "=="))
+NUMPY_NAN = frozenset(("nan", "NaN", "NAN"))
+
+
+class NumpyNaNComparisonChecker(BaseChecker):
+    name = "numpy-nan-compare"
+    msgs = {
+        "W8001": (
+            "Numpy nan comparison used",
+            "numpy-nan-compare",
+            "Since comparing NaN with NaN always returns False, use np.isnan() to check for NaN values.",
+        ),
+    }
+
+    @classmethod
+    def __is_np_nan_call(cls, node):
+        """Check if the node represents a call to np.nan."""
+        return (
+            isinstance(node, nodes.Call)
+            and isinstance(node.func, nodes.Attribute)
+            and node.func.attrname in NUMPY_NAN
+            and isinstance(node.func.expr, nodes.Name)
+            and node.func.expr.name == "np"
+        )
+
+    @only_required_for_messages("numpy-nan-compare")
+    def visit_compare(self, node: nodes.Compare) -> None:
+        # Why am I getting node: node.Call here? Should be nodes.Compare...
+        # Check test case test_numpy/test_numpy_nan_comparison.py
+        if self.__is_np_nan_call(node.left):
+            self.add_message("numpy-nan-compare", node=node, confidence=HIGH)
+            return
+
+        for op, comparator in node.ops:
+            if op in COMPARISON_OP and self.__is_np_nan_call(comparator):
+                self.add_message("numpy-nan-compare", node=node, confidence=HIGH)
+                return

pylint_ml/checkers/numpy/numpy_nan_comparison.py

@@ -2,18 +2,20 @@
 
 from pylint.lint import PyLinter
 
-from pylint_ml.checkers.numpy.import_numpy import NumpyImportChecker
-from pylint_ml.checkers.pandas.import_pandas import PandasImportChecker
-from pylint_ml.checkers.scipy.import_scipy import ScipyImportChecker
-from pylint_ml.checkers.sklearn.import_sklearn import SklearnImportChecker
-from pylint_ml.checkers.tensorflow.import_tensorflow import TensorflowImportChecker
-from pylint_ml.checkers.torch.import_torch import TorchImportChecker
+from pylint_ml.checkers.numpy.numpy_import import NumpyImportChecker
+from pylint_ml.checkers.numpy.numpy_nan_comparison import NumpyNaNComparisonChecker
+from pylint_ml.checkers.pandas.pandas_import import PandasImportChecker
+from pylint_ml.checkers.scipy.scipy_import import ScipyImportChecker
+from pylint_ml.checkers.sklearn.sklearn_import import SklearnImportChecker
+from pylint_ml.checkers.tensorflow.tensorflow_import import TensorflowImportChecker
+from pylint_ml.checkers.torch.torch_import import TorchImportChecker
 
 
 def register(linter: PyLinter) -> None:
     """Register checkers."""
     # Numpy
     linter.register_checker(NumpyImportChecker(linter))
+    linter.register_checker(NumpyNaNComparisonChecker(linter))
 
     # Pandas
     linter.register_checker(PandasImportChecker(linter))

pylint_ml/checkers/pandas/import_pandas.py renamed to pylint_ml/checkers/pandas/pandas_import.py

@@ -2,7 +2,7 @@
 import pylint.testutils
 from pylint.interfaces import HIGH
 
-from pylint_ml.checkers.numpy.import_numpy import NumpyImportChecker
+from pylint_ml.checkers.numpy.numpy_import import NumpyImportChecker
 
 
 class TestNumpyImport(pylint.testutils.CheckerTestCase):

pylint_ml/checkers/scipy/import_scipy.py renamed to pylint_ml/checkers/scipy/scipy_import.py

@@ -0,0 +1,47 @@
+import astroid
+import pylint.testutils
+from pylint.interfaces import HIGH
+
+from pylint_ml.checkers.numpy.numpy_nan_comparison import NumpyNaNComparisonChecker
+
+
+class TestNumpyNaNComparison(pylint.testutils.CheckerTestCase):
+    CHECKER_CLASS = NumpyNaNComparisonChecker
+
+    def test_correct_nan_compare(self):
+        nan_compare_node = astroid.extract_node(
+            """
+        np.isnan(np.nan)
+        """
+        )
+
+        with self.assertNoMessages():
+            self.checker.visit_compare(nan_compare_node)
+
+    def test_incorrect_nan_compare(self):
+        nan_compare_node = astroid.extract_node(
+            """
+        a_nan = np.array([0, 1, np.nan])
+        print(a_nan)
+        # [ 0.  1. nan]
+
+        print(a_nan == np.nan)
+        # [False False False]
+
+        print(np.isnan(a_nan))
+        # [False False  True]
+
+        print(a_nan > 0)
+        # [False  True False]
+        """
+        )
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="numpy-nan-compare",
+                confidence=HIGH,
+                node=nan_compare_node,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_compare(nan_compare_node)