Merge pull request #13 from pylint-dev/12-add-numpy-nan-comparision-checker

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,43 @@
+# 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: nodes.Attribute) -> bool:
+        """Check if the node represents a call to np.nan."""
+        return node.attrname in NUMPY_NAN and isinstance(node.expr, nodes.Name) and node.expr.name == "np"
+
+    @only_required_for_messages("numpy-nan-compare")
+    def visit_compare(self, node: nodes.Compare) -> None:
+
+        if isinstance(node.left, nodes.Attribute) and 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 isinstance(comparator, nodes.Attribute) 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,44 @@
+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_singleton_nan_compare(self):
+        code = """
+        a_nan = np.array([0, 1, np.nan])
+
+        np.nan == a_nan #@
+
+        1 == 1 == np.nan #@
+
+        1 > 0 > np.nan #@
+
+        """
+        singleton_nan_compare, chained_nan_compare, great_than_nan_compare = astroid.extract_node(code)
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="numpy-nan-compare",
+                node=singleton_nan_compare,
+                confidence=HIGH,
+            ),
+            pylint.testutils.MessageTest(
+                msg_id="numpy-nan-compare",
+                node=chained_nan_compare,
+                confidence=HIGH,
+            ),
+            pylint.testutils.MessageTest(
+                msg_id="numpy-nan-compare",
+                node=great_than_nan_compare,
+                confidence=HIGH,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_compare(singleton_nan_compare)
+            self.checker.visit_compare(chained_nan_compare)
+            self.checker.visit_compare(great_than_nan_compare)