From e8539d208f1254378539cf28b59161e0872be5ee Mon Sep 17 00:00:00 2001
From: Marco Gorelli <marcogorelli@protonmail.com>
Date: Mon, 11 Oct 2021 14:16:09 +0100
Subject: [PATCH] add check for array_equal and np.testing

---
 README.md                                     |  2 +
 .../_plugins_tree/numpy_testing.py            | 40 ++++++++++
 tests/numpy_testing_test.py                   | 77 +++++++++++++++++++
 3 files changed, 119 insertions(+)
 create mode 100644 pandas_dev_flaker/_plugins_tree/numpy_testing.py
 create mode 100644 tests/numpy_testing_test.py

diff --git a/README.md b/README.md
index 644a61a..bb649e7 100644
--- a/README.md
+++ b/README.md
@@ -41,6 +41,8 @@ a linter for pandas usage, please see [pandas-vet](https://github.com/deppen8/pa
 | PDF022 | found import from 'numpy.random'                                        |
 | PDF023 | found assignment to single-letter variable                              |
 | PDF024 | found string join() with generator expressions                          |
+| PDF025 | found 'np.testing' or 'np.array_equal' (use 'pandas._testing' instead)  |
+
 ## contributing
 
 See `contributing.md` for how to get started.
diff --git a/pandas_dev_flaker/_plugins_tree/numpy_testing.py b/pandas_dev_flaker/_plugins_tree/numpy_testing.py
new file mode 100644
index 0000000..ef56468
--- /dev/null
+++ b/pandas_dev_flaker/_plugins_tree/numpy_testing.py
@@ -0,0 +1,40 @@
+import ast
+from typing import Iterator, Tuple
+
+from pandas_dev_flaker._data_tree import State, register
+
+MSG = (
+    "PDF025 found 'np.testing' or 'np.array_equal' "
+    "(use 'pandas._testing' instead)"
+)
+
+
+@register(ast.ImportFrom)
+def visit_ImportFrom(
+    state: State,
+    node: ast.ImportFrom,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
+    if (
+        (
+            "testing" in {name.name for name in node.names}
+            or "array_equal" in {name.name for name in node.names}
+        )
+        and node.module is not None
+        and node.module == "numpy"
+    ) or (node.module is not None and node.module == "numpy.testing"):
+        yield node.lineno, node.col_offset, MSG
+
+
+@register(ast.Attribute)
+def visit_Attribute(
+    state: State,
+    node: ast.Attribute,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
+    if (
+        node.attr in {"testing", "array_equal"}
+        and isinstance(node.value, ast.Name)
+        and node.value.id in {"numpy", "np"}
+    ):
+        yield node.lineno, node.col_offset, MSG
diff --git a/tests/numpy_testing_test.py b/tests/numpy_testing_test.py
new file mode 100644
index 0000000..363ad66
--- /dev/null
+++ b/tests/numpy_testing_test.py
@@ -0,0 +1,77 @@
+import ast
+import tokenize
+from io import StringIO
+
+import pytest
+
+from pandas_dev_flaker.__main__ import run
+
+
+def results(s):
+    return {
+        "{}:{}: {}".format(*r)
+        for r in run(
+            ast.parse(s),
+            list(tokenize.generate_tokens(StringIO(s).readline)),
+        )
+    }
+
+
+@pytest.mark.parametrize(
+    "source",
+    (
+        pytest.param(
+            "import numpy as np\nnp.tester",
+            id="tester instead of testing",
+        ),
+    ),
+)
+def test_noop(source):
+    assert not results(source)
+
+
+@pytest.mark.parametrize(
+    "source, expected",
+    (
+        pytest.param(
+            "import numpy\nnumpy.testing",
+            "2:0: PDF025 found 'np.testing' or 'np.array_equal' "
+            "(use 'pandas._testing' instead)",
+            id="access from numpy",
+        ),
+        pytest.param(
+            "import numpy as np\nnp.testing",
+            "2:0: PDF025 found 'np.testing' or 'np.array_equal' "
+            "(use 'pandas._testing' instead)",
+            id="access from np",
+        ),
+        pytest.param(
+            "import numpy as np\n"
+            "np.testing.assert_array_equal(result, expected)",
+            "2:0: PDF025 found 'np.testing' or 'np.array_equal' "
+            "(use 'pandas._testing' instead)",
+            id="np.testing.assert_array_equal",
+        ),
+        pytest.param(
+            "from numpy import testing",
+            "1:0: PDF025 found 'np.testing' or 'np.array_equal' "
+            "(use 'pandas._testing' instead)",
+            id="import from numpy",
+        ),
+        pytest.param(
+            "from numpy.testing import assert_array_equal",
+            "1:0: PDF025 found 'np.testing' or 'np.array_equal' "
+            "(use 'pandas._testing' instead)",
+            id="import from numpy",
+        ),
+        pytest.param(
+            "from numpy import array_equal",
+            "1:0: PDF025 found 'np.testing' or 'np.array_equal' "
+            "(use 'pandas._testing' instead)",
+            id="array_equal",
+        ),
+    ),
+)
+def test_violation(source, expected):
+    (result,) = results(source)
+    assert result == expected