Fixup collections

.pre-commit-config.yaml

@@ -18,6 +18,7 @@ repos:
     hooks:
     -   id: flake8
         additional_dependencies: [flake8-typing-imports==1.7.0]
+        args: [--extend-ignore=E203]
 -   repo: https://github.com/psf/black
     rev: 20.8b1
     hooks:

pandas_dev_flaker/_plugins_tokens/strings_to_concatenate.py

@@ -8,7 +8,7 @@
 
 
 @register()
-def foo(
+def visit(
     tokens: Sequence[tokenize.TokenInfo],
 ) -> Iterator[Tuple[int, int, str]]:
     for current_token, next_token in zip(tokens, tokens[1:]):

pandas_dev_flaker/_plugins_tree/collections_abc.py

@@ -3,7 +3,10 @@
 
 from pandas_dev_flaker._data_tree import State, register
 
-MSG = "PDF001 don't import from collections.abc"
+MSG = (
+    "PDF001 import from collections.abc "
+    "(use 'from collections import abc' instead)"
+)
 
 
 @register(ast.ImportFrom)
@@ -14,3 +17,13 @@ def visit_ImportFrom(
 ) -> Iterator[Tuple[int, int, str]]:
     if node.module == "collections.abc":
         yield node.lineno, node.col_offset, MSG
+
+
+@register(ast.Import)
+def visit_Import(
+    state: State,
+    node: ast.Import,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
+    if "collections.abc" in {name.name for name in node.names}:
+        yield node.lineno, node.col_offset, MSG

tests/collections_abc_test.py

@@ -0,0 +1,52 @@
+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(
+            "from collections import abc",
+            id="Imported abc from collections",
+        ),
+    ),
+)
+def test_noop(source):
+    assert not results(source)
+
+
+@pytest.mark.parametrize(
+    "source, expected",
+    (
+        pytest.param(
+            "from collections.abc import Generator",
+            "1:0: PDF001 import from collections.abc "
+            "(use 'from collections import abc' instead)",
+            id="Imported from collections.abc",
+        ),
+        pytest.param(
+            "import collections.abc",
+            "1:0: PDF001 import from collections.abc "
+            "(use 'from collections import abc' instead)",
+            id="Imported collections.abc",
+        ),
+    ),
+)
+def test_violation(source, expected):
+    (result,) = results(source)
+    assert result == expected