Merge pull request #26 from vrserpa/generator_expressions_join_strings_check

Linting check of use of string join() with generator expression

Author: MarcoGorelli

README.md

@@ -40,6 +40,7 @@ a linter for pandas usage, please see [pandas-vet](https://github.com/deppen8/pa
 | PDF021 | found 'np.bool' or 'np.object' (use 'np.bool_' or 'np.object_' instead) |
 | PDF022 | found import from 'numpy.random'                                        |
 | PDF023 | found assignment to single-letter variable                              |
+| PDF024 | found string join() with generator expressions                          |
 ## contributing
 
 See `contributing.md` for how to get started.

pandas_dev_flaker/_ast_helpers.py

@@ -1,4 +1,5 @@
 import ast
+import sys
 from typing import Container, Dict, Sequence, Set
 
 
@@ -30,3 +31,19 @@ def check_for_wrong_alias(
             return name_.asname != alias
     else:
         return False
+
+
+def is_str_constant(
+    node: ast.Call,
+) -> bool:
+    return isinstance(node.func, ast.Attribute) and (
+        (
+            sys.version_info < (3, 8)
+            and isinstance(node.func.value, ast.Str)
+        )
+        or (
+            sys.version_info >= (3, 8)
+            and isinstance(node.func.value, ast.Constant)
+            and isinstance(node.func.value.value, str)
+        )
+    )

pandas_dev_flaker/_plugins_tree/generator_join_strings.py

@@ -0,0 +1,29 @@
+"""
+Based on
+
+https://github.com/asottile/pyupgrade/blob/5fb168667ae73f157dd579344708e1cdfb0c0341/pyupgrade/_plugins/generator_expressions_pep289.py
+"""
+
+import ast
+from typing import Iterator, Tuple
+
+from pandas_dev_flaker._ast_helpers import is_str_constant
+from pandas_dev_flaker._data_tree import State, register
+
+MSG = "PDF024 found string join() with generator expressions"
+
+
+@register(ast.Call)
+def visit_Call(
+    state: State,
+    node: ast.Call,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
+    if (
+        isinstance(node.func, ast.Attribute)
+        and node.func.attr == "join"
+        and is_str_constant(node)
+        and node.args
+        and isinstance(node.args[0], ast.GeneratorExp)
+    ):
+        yield (node.lineno, node.col_offset, MSG)

tests/generator_join_strings_test.py

@@ -0,0 +1,71 @@
+"""
+Based on
+
+https://github.com/asottile/pyupgrade/blob/5fb168667ae73f157dd579344708e1cdfb0c0341/pyupgrade/_plugins/generator_expressions_pep289.py
+"""
+
+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(
+            "''.join([str(i) for i in range(5)])",
+            id="String join() with list comprehension",
+        ),
+        pytest.param(
+            "''.join([\n"
+            "        str(i) for i in range(5)\n"
+            "        ]\n"
+            ")\n",
+            id="String join() with multiline list comprehension",
+        ),
+    ),
+)
+def test_list_comprehensions(source):
+    assert not results(source)
+
+
+@pytest.mark.parametrize(
+    "source, expected",
+    (
+        pytest.param(
+            "''.join(str(i) for i in range(5))",
+            "1:0: PDF024 found string join() with generator expressions",
+            id="String join() with generator expression",
+        ),
+        pytest.param(
+            "''.join((str(i) for i in range(5)))",
+            "1:0: PDF024 found string join() with generator expressions",
+            id="String join() with parenthesised generator expression",
+        ),
+        pytest.param(
+            "''.join((\n"
+            "        str(i) for i in range(5)\n"
+            "        )\n"
+            ")\n",
+            "1:0: PDF024 found string join() with generator expressions",
+            id="String join() with multiline generator expression",
+        ),
+    ),
+)
+def test_generator_expressions(source, expected):
+    (result,) = results(source)
+    assert result == expected