Add nan compare checker final

Author: PeterHamfelt

pylint_ml/checkers/numpy/numpy_nan_comparison.py

@@ -26,25 +26,18 @@ class NumpyNaNComparisonChecker(BaseChecker):
     }
 
     @classmethod
-    def __is_np_nan_call(cls, node):
+    def __is_np_nan_call(cls, node: nodes.Attribute) -> bool:
         """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"
-        )
+        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:
-        # 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):
+
+        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 self.__is_np_nan_call(comparator):
+            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

tests/checkers/test_numpy/test_numpy_nan_comparison.py

@@ -8,40 +8,37 @@
 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(
-            """
+    def test_singleton_nan_compare(self):
+        code = """
         a_nan = np.array([0, 1, np.nan])
-        print(a_nan)
-        # [ 0.  1. nan]
 
-        print(a_nan == np.nan)
-        # [False False False]
+        np.nan == a_nan #@
 
-        print(np.isnan(a_nan))
-        # [False False  True]
+        1 == 1 == np.nan #@
+
+        1 > 0 > np.nan #@
 
-        print(a_nan > 0)
-        # [False  True False]
         """
-        )
+        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,
-                node=nan_compare_node,
             ),
             ignore_position=True,
         ):
-            self.checker.visit_compare(nan_compare_node)
+            self.checker.visit_compare(singleton_nan_compare)
+            self.checker.visit_compare(chained_nan_compare)
+            self.checker.visit_compare(great_than_nan_compare)