added checks for return type annotations and variable assignment annotations

Author: Pydare

pandas_dev_flaker/_plugins_tree/disallow_argument_types.py

@@ -6,21 +6,57 @@
 MSG = "PDF026 found union between Series and AnyArrayLike in type hint"
 
 
+# for function arguments annotations
 @register(ast.FunctionDef)
-def visit_arg(
+def visit_FunctionDef(
     state: State,
     node: ast.FunctionDef,
     parent: ast.AST,
 ) -> Iterator[Tuple[int, int, str]]:
     arguments = node.args.args
     for arg in arguments:
-        if isinstance(arg.annotation, ast.BinOp) and parse_annotation(
+        if isinstance(
+            arg.annotation,
+            ast.BinOp,
+        ) and _contains_series_and_arraylike(
             arg.annotation,
         ):
             yield arg.lineno, arg.col_offset, MSG
 
 
-def parse_annotation(node: ast.BinOp) -> bool:
+# for return arguments annotations
+@register(ast.FunctionDef)
+def visit_args(
+    state: State,
+    node: ast.FunctionDef,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
+    return_arguments = node.returns  # BinOp, Name, Subscript
+    if (
+        isinstance(
+            return_arguments,
+            ast.BinOp,
+        )
+        and _contains_series_and_arraylike(return_arguments)
+    ):
+        yield node.lineno, node.col_offset, MSG
+
+
+# for annotations defined outside function args & return args
+@register(ast.AnnAssign)
+def visit_AnnAssign(
+    state: State,
+    node: ast.AnnAssign,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
+    annotations = node.annotation
+    if isinstance(annotations, ast.BinOp) and _contains_series_and_arraylike(
+        annotations,
+    ):
+        yield node.lineno, node.col_offset, MSG
+
+
+def _contains_series_and_arraylike(node: ast.BinOp) -> bool:
     series, any_array_like = "Series", "AnyArrayLike"
     is_series, is_array_like = False, False
 

tests/disallow_argument_types_test.py

@@ -49,6 +49,10 @@ def results(s):
             "def f(a: Callable[..., T] | DataFrame | list[int]): pass",
             id="Function annotation containing Subscript type",
         ),
+        pytest.param(
+            "def f(a: DataFrame | list[int]) -> int | str: pass",
+            id="Function return annotation containing Subscript type",
+        ),
     ),
 )
 def test_noop(source):
@@ -96,3 +100,97 @@ def test_noop(source):
 def test_violation(source, expected):
     (result,) = results(source)
     assert result == expected
+
+
+@pytest.mark.parametrize(
+    "source",
+    (
+        pytest.param(
+            "def f(foo) -> int | str | bool: pass",
+            id="Function with multiple return type annotations",
+        ),
+        pytest.param(
+            "def foo(bar: list[int]): pass",
+            id="Function with no return type",
+        ),
+        pytest.param(
+            "def foo(self, bar: int) -> int: pass",
+            id="Function with one return type annotation",
+        ),
+    ),
+)
+def test_noop2(source):
+    assert not results(source)
+
+
+@pytest.mark.parametrize(
+    "source, expected",
+    (
+        pytest.param(
+            "def bar(foo, other: tuple[Callable[..., T]] | "
+            "Series | list[int]) -> Series | AnyArrayLike | "
+            "DataFrame: pass",
+            "1:0: PDF026 found union between Series and "
+            "AnyArrayLike in "
+            "type hint",
+            id="found union between Series and AnyArrayLike "
+            "in return annotations",
+        ),
+        pytest.param(
+            "def bar(foo: int, other: tuple[Callable[..., T]] | "
+            "Series | list[int]) -> Series | AnyArrayLike: pass",
+            "1:0: PDF026 found union between Series and "
+            "AnyArrayLike in "
+            "type hint",
+            id="found union between Series and AnyArrayLike "
+            "in return annotations",
+        ),
+    ),
+)
+def test_violation2(source, expected):
+    (result,) = results(source)
+    assert result == expected
+
+
+@pytest.mark.parametrize(
+    "source",
+    (
+        pytest.param(
+            "foo: str = 'string variable'",
+            id="Assignment with one annotation",
+        ),
+        pytest.param(
+            "self.bar: DataFrame | Timezone = [1, 2, 3]",
+            id="Assignment with multiple annotations",
+        ),
+        pytest.param("cls.foo = 3", id="Assignment with no annotation"),
+    ),
+)
+def test_noop3(source):
+    assert not results(source)
+
+
+@pytest.mark.parametrize(
+    "source, expected",
+    (
+        pytest.param(
+            "self.foo: AnyArrayLike | Timezone | Series = 2",
+            "1:0: PDF026 found union between Series and "
+            "AnyArrayLike in "
+            "type hint",
+            id="found union between Series and AnyArrayLike "
+            "in variable assignment",
+        ),
+        pytest.param(
+            "cls.foo: AnyArrayLike | Series = 2",
+            "1:0: PDF026 found union between Series and "
+            "AnyArrayLike in "
+            "type hint",
+            id="found union between Series and AnyArrayLike "
+            "in variable assignment",
+        ),
+    ),
+)
+def test_violation3(source, expected):
+    (result,) = results(source)
+    assert result == expected