diff --git a/README.md b/README.md index 5ec40ee..6d51ced 100644 --- a/README.md +++ b/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. diff --git a/pandas_dev_flaker/_ast_helpers.py b/pandas_dev_flaker/_ast_helpers.py index 168e1b9..1558f5d 100644 --- a/pandas_dev_flaker/_ast_helpers.py +++ b/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) + ) + ) diff --git a/pandas_dev_flaker/_plugins_tree/generator_join_strings.py b/pandas_dev_flaker/_plugins_tree/generator_join_strings.py new file mode 100644 index 0000000..346ff8d --- /dev/null +++ b/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) diff --git a/tests/generator_join_strings_test.py b/tests/generator_join_strings_test.py new file mode 100644 index 0000000..20e4c0f --- /dev/null +++ b/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