Skip to content
This repository was archived by the owner on Feb 19, 2023. It is now read-only.

Commit c2586bb

Browse files
MarcoGorelliMarcoGorelli
authored
Merge pull request #28 from Pydare/disallow_series_arraylike
disallowed Series | AnyArrayLike as Series is already array-like
2 parents 4bd7b6e + 185fd23 commit c2586bb

File tree

4 files changed

+261
-4
lines changed

4 files changed

+261
-4
lines changed

‎README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ a linter for pandas usage, please see [pandas-vet](https://github.com/deppen8/pa
4242
| PDF023 | found assignment to single-letter variable |
4343
| PDF024 | found string join() with generator expressions |
4444
| PDF025 | found 'np.testing' or 'np.array_equal' (use 'pandas._testing' instead) |
45+
| PDF026 | found union between Series and AnyArrayLike in type hint |
4546

4647
## contributing
4748

‎pandas_dev_flaker/_ast_helpers.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,10 +37,7 @@ def is_str_constant(
3737
node: ast.Call,
3838
) -> bool:
3939
return isinstance(node.func, ast.Attribute) and (
40-
(
41-
sys.version_info < (3, 8)
42-
and isinstance(node.func.value, ast.Str)
43-
)
40+
(sys.version_info < (3, 8) and isinstance(node.func.value, ast.Str))
4441
or (
4542
sys.version_info >= (3, 8)
4643
and isinstance(node.func.value, ast.Constant)

‎pandas_dev_flaker/_plugins_tree/disallow_argument_types.py

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
import ast
2+
from typing import Iterator, Tuple
3+
4+
from pandas_dev_flaker._data_tree import State, register
5+
6+
MSG = "PDF026 found union between Series and AnyArrayLike in type hint"
7+
SERIES, ANY_ARRAY_LIKE = "Series", "AnyArrayLike"
8+
9+
10+
def _contains_series_and_arraylike(node: ast.AST) -> bool:
11+
ret = False
12+
for node in ast.walk(node):
13+
if isinstance(node, ast.BinOp):
14+
ret |= _binop_contains_series_and_arraylike(node)
15+
return ret
16+
17+
18+
def _binop_contains_series_and_arraylike(node: ast.BinOp) -> bool:
19+
is_series, is_array_like = False, False
20+
21+
for _node in ast.walk(node):
22+
if isinstance(_node, ast.Name):
23+
if _node.id == SERIES:
24+
is_series = True
25+
elif _node.id == ANY_ARRAY_LIKE:
26+
is_array_like = True
27+
elif isinstance(_node, ast.Str):
28+
if _node.s == SERIES:
29+
is_series = True
30+
elif _node.s == ANY_ARRAY_LIKE:
31+
is_array_like = True
32+
33+
return is_series and is_array_like
34+
35+
36+
# for function arguments/returns annotations
37+
@register(ast.FunctionDef)
38+
def visit_FunctionDef(
39+
state: State,
40+
node: ast.FunctionDef,
41+
parent: ast.AST,
42+
) -> Iterator[Tuple[int, int, str]]:
43+
arguments = node.args.args
44+
for arg in arguments:
45+
if arg.annotation is not None and _contains_series_and_arraylike(
46+
arg.annotation,
47+
):
48+
yield arg.lineno, arg.col_offset, MSG
49+
if node.returns is not None and _contains_series_and_arraylike(
50+
node.returns,
51+
):
52+
yield node.lineno, node.col_offset, MSG
53+
54+
55+
# for annotations defined outside function args & return args
56+
@register(ast.AnnAssign)
57+
def visit_AnnAssign(
58+
state: State,
59+
node: ast.AnnAssign,
60+
parent: ast.AST,
61+
) -> Iterator[Tuple[int, int, str]]:
62+
annotation = node.annotation
63+
if annotation is not None and _contains_series_and_arraylike(annotation):
64+
yield node.lineno, node.col_offset, MSG

‎tests/disallow_argument_types_test.py

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
import ast
2+
import tokenize
3+
from io import StringIO
4+
5+
import pytest
6+
7+
from pandas_dev_flaker.__main__ import run
8+
9+
10+
def results(s):
11+
return {
12+
"{}:{}: {}".format(*r)
13+
for r in run(
14+
ast.parse(s),
15+
list(tokenize.generate_tokens(StringIO(s).readline)),
16+
)
17+
}
18+
19+
20+
@pytest.mark.parametrize(
21+
"source",
22+
(
23+
pytest.param(
24+
"def f(foo): pass",
25+
id="Function argument with no annotation ",
26+
),
27+
pytest.param(
28+
"def f(foo, other: Series): pass",
29+
id="Function argument with one annotation ",
30+
),
31+
pytest.param(
32+
"def p(foo, other: Series | DataFrame ): pass",
33+
id="Function argument with two annotations",
34+
),
35+
pytest.param(
36+
"def p(foo, other: Series | Union[int, str] ): pass",
37+
id="Function argument with two annotations",
38+
),
39+
pytest.param(
40+
"def q(foo, other: DataFrame | AnyArrayLike | Timestamp): pass",
41+
id="Function argument with three annotations" "AnyArrayLike",
42+
),
43+
pytest.param(
44+
"def b(foo, other: DataFrame | Timezone | "
45+
"Timestamp | Timedelta): pass",
46+
id="Function argument with four annotations",
47+
),
48+
pytest.param(
49+
"def f(a: Callable[..., T] | DataFrame | list[int]): pass",
50+
id="Function annotation containing Subscript type",
51+
),
52+
pytest.param(
53+
"def f(a: DataFrame | list[int]) -> int | str: pass",
54+
id="Function return annotation containing Subscript type",
55+
),
56+
),
57+
)
58+
def test_noop(source):
59+
assert not results(source)
60+
61+
62+
@pytest.mark.parametrize(
63+
"source, expected",
64+
(
65+
pytest.param(
66+
"def dot(foo, other: AnyArrayLike | Series): pass",
67+
"1:13: PDF026 found union between Series and "
68+
"AnyArrayLike in "
69+
"type hint",
70+
id="Series and AnyArrayLike",
71+
),
72+
pytest.param(
73+
"def bar(foo, other: DataFrame | Series | AnyArrayLike): pass",
74+
"1:13: PDF026 found union between Series and "
75+
"AnyArrayLike in "
76+
"type hint",
77+
id="Series and AnyArrayLike " "and one other annotation",
78+
),
79+
pytest.param(
80+
"def bar(foo, other: DataFrame | Series | "
81+
"AnyArrayLike | NDFrame): pass",
82+
"1:13: PDF026 found union between Series and "
83+
"AnyArrayLike in "
84+
"type hint",
85+
id="Series and AnyArrayLike " "and two other annotations",
86+
),
87+
),
88+
)
89+
def test_violation(source, expected):
90+
(result,) = results(source)
91+
assert result == expected
92+
93+
94+
@pytest.mark.parametrize(
95+
"source",
96+
(
97+
pytest.param(
98+
"def f(foo) -> int | str | bool: pass",
99+
id="Function with multiple return type annotations",
100+
),
101+
pytest.param(
102+
"def foo(bar: list[int]): pass",
103+
id="Function with no return type",
104+
),
105+
pytest.param(
106+
"def foo(self, bar: int) -> int: pass",
107+
id="Function with one return type annotation",
108+
),
109+
),
110+
)
111+
def test_noop_returns(source):
112+
assert not results(source)
113+
114+
115+
@pytest.mark.parametrize(
116+
"source, expected",
117+
(
118+
pytest.param(
119+
"def bar(foo, other: tuple[Callable[..., T]] | "
120+
"Series | list[int]) -> Series | AnyArrayLike | "
121+
"DataFrame: pass",
122+
"1:0: PDF026 found union between Series and "
123+
"AnyArrayLike in "
124+
"type hint",
125+
id="3 objects in return type",
126+
),
127+
pytest.param(
128+
"def bar(foo: int, other: tuple[Callable[..., T]] | "
129+
"Series | list[int]) -> Series | AnyArrayLike: pass",
130+
"1:0: PDF026 found union between Series and "
131+
"AnyArrayLike in "
132+
"type hint",
133+
id="2 objects in return type",
134+
),
135+
pytest.param(
136+
"def bar(foo: List[Series | AnyArrayLike]): ...",
137+
"1:8: PDF026 found union between Series and "
138+
"AnyArrayLike in "
139+
"type hint",
140+
id="List of Series or AnyArrayLike",
141+
),
142+
pytest.param(
143+
"def bar(foo: List['Series' | 'AnyArrayLike' | 'int']): ...",
144+
"1:8: PDF026 found union between Series and "
145+
"AnyArrayLike in "
146+
"type hint",
147+
id="String version of Series",
148+
),
149+
),
150+
)
151+
def test_violation_returns(source, expected):
152+
(result,) = results(source)
153+
assert result == expected
154+
155+
156+
@pytest.mark.parametrize(
157+
"source",
158+
(
159+
pytest.param(
160+
"foo: str = 'string variable'",
161+
id="Assignment with one annotation",
162+
),
163+
pytest.param(
164+
"self.bar: DataFrame | Timezone = [1, 2, 3]",
165+
id="Assignment with multiple annotations",
166+
),
167+
pytest.param("cls.foo = 3", id="Assignment with no annotation"),
168+
),
169+
)
170+
def test_noop_assignment(source):
171+
assert not results(source)
172+
173+
174+
@pytest.mark.parametrize(
175+
"source, expected",
176+
(
177+
pytest.param(
178+
"self.foo: AnyArrayLike | Timezone | Series = 2",
179+
"1:0: PDF026 found union between Series and "
180+
"AnyArrayLike in "
181+
"type hint",
182+
id="annotation with assignment",
183+
),
184+
pytest.param(
185+
"foo: AnyArrayLike | Series",
186+
"1:0: PDF026 found union between Series and "
187+
"AnyArrayLike in "
188+
"type hint",
189+
id="simple annotation",
190+
),
191+
),
192+
)
193+
def test_violation_assignment(source, expected):
194+
(result,) = results(source)
195+
assert result == expected

0 commit comments

Comments
 (0)