wip

Author: MarcoGorelli

Diff for: pandas_style_guide/__main__.py

@@ -1,18 +1,13 @@
 from __future__ import annotations
-import tokenize
-import os
 
 import ast
 import importlib.metadata
-from typing import Any, List, Dict, Tuple, Type, Callable
-import collections
-from typing import Generator
+import os
+from typing import Any, Generator
 
 from pandas_style_guide._data import FUNCS, visit
 
 
-
-
 class Plugin:
     name = os.path.split(os.path.dirname(__file__))[-1]
     version = importlib.metadata.version(name)
@@ -26,4 +21,3 @@ def run(self) -> Generator[tuple[int, int, str, type[Any]], None, None]:
             return
         for line, col, msg in callbacks:
             yield line, col, msg, type(self)
-        

Diff for: pandas_style_guide/_ast_helpers.py

@@ -1,42 +1,34 @@
 import ast
 import warnings
-from typing import Container
-from typing import Dict
-from typing import Set
-from typing import Union
-
+from typing import Container, Dict, Set
 
 
 def ast_parse(contents_text: str) -> ast.Module:
     # intentionally ignore warnings, we might be fixing warning-ridden syntax
     with warnings.catch_warnings():
-        warnings.simplefilter('ignore')
+        warnings.simplefilter("ignore")
         return ast.parse(contents_text.encode())
 
 
-
-
 def is_name_attr(
-        node: ast.AST,
-        imports: Dict[str, Set[str]],
-        mod: str,
-        names: Container[str],
+    node: ast.AST,
+    imports: Dict[str, Set[str]],
+    mod: str,
+    names: Container[str],
 ) -> bool:
     return (
-        isinstance(node, ast.Name) and
-        node.id in names and
-        node.id in imports[mod]
+        isinstance(node, ast.Name)
+        and node.id in names
+        and node.id in imports[mod]
     ) or (
-        isinstance(node, ast.Attribute) and
-        isinstance(node.value, ast.Name) and
-        node.value.id == mod and
-        node.attr in names
+        isinstance(node, ast.Attribute)
+        and isinstance(node.value, ast.Name)
+        and node.value.id == mod
+        and node.attr in names
     )
 
 
 def has_starargs(call: ast.Call) -> bool:
-    return (
-        any(k.arg is None for k in call.keywords) or
-        any(isinstance(a, ast.Starred) for a in call.args)
+    return any(k.arg is None for k in call.keywords) or any(
+        isinstance(a, ast.Starred) for a in call.args
     )
-

Diff for: pandas_style_guide/_data.py

@@ -1,48 +1,72 @@
-import pkgutil
-import collections
-from typing import List, Tuple, Dict, NamedTuple, Set
 import ast
+import collections
+import pkgutil
+from typing import (
+    Dict,
+    Iterator,
+    List,
+    NamedTuple,
+    Protocol,
+    Sequence,
+    Set,
+    Tuple,
+    Type,
+    TypeVar,
+)
+
 from pandas_style_guide import _plugins
 
 FUNCS = collections.defaultdict(list)
+AST_T = TypeVar("AST_T", bound=ast.AST)
+
+
+class ASTCallbackMapping(Protocol):
+    def __getitem__(self, tp: Type[AST_T]) -> Tuple[int, int, str]:
+        ...
+
+
 class State(NamedTuple):
     from_imports: Dict[str, Set[str]]
     in_annotation: bool = False
 
+
 def register(tp):
     def register_decorator(func):
         FUNCS[tp].append(func)
         return func
+
     return register_decorator
 
+
 def _get_alias(name):
     if name.asname is not None:
         return name.asname
     else:
         return name.name
 
-def visit(funcs, tree: ast.Module) -> Dict[int, List[int]]:
+
+def visit(
+    funcs: ASTCallbackMapping,
+    tree: ast.Module,
+) -> Iterator[Tuple[str, str, int]]:
     "Step through tree, recording when nodes are in annotations."
     initial_state = State(
         from_imports=collections.defaultdict(set),
     )
-    nodes: List[Tuple[bool, ast.AST, ast.AST]] = [(initial_state, tree, tree)]
+    nodes: List[Tuple[State, ast.AST, ast.AST]] = [(initial_state, tree, tree)]
 
     while nodes:
         state, node, parent = nodes.pop()
         tp = type(node)
         for ast_func in funcs[tp]:
             yield from ast_func(state, node, parent)
 
-        if (
-                isinstance(node, ast.ImportFrom)
-                and node.module is not None
-        ):
-            state.from_imports[node.module.split('.')[0]].update(
+        if isinstance(node, ast.ImportFrom) and node.module is not None:
+            state.from_imports[node.module.split(".")[0]].update(
                 _get_alias(name) for name in node.names if not name.asname
             )
         elif isinstance(node, ast.Import):
-            for name in node.names: 
+            for name in node.names:
                 state.from_imports[_get_alias(name)]
 
         for name in reversed(node._fields):
@@ -58,12 +82,13 @@ def visit(funcs, tree: ast.Module) -> Dict[int, List[int]]:
                     if isinstance(value, ast.AST):
                         nodes.append((next_state, value, node))
 
+
 def _import_plugins() -> None:
     # https://github.com/python/mypy/issues/1422
     plugins_path: str = _plugins.__path__  # type: ignore
-    mod_infos = pkgutil.walk_packages(plugins_path, f'{_plugins.__name__}.')
+    mod_infos = pkgutil.walk_packages(plugins_path, f"{_plugins.__name__}.")
     for _, name, _ in mod_infos:
-        __import__(name, fromlist=['_trash'])
+        __import__(name, fromlist=["_trash"])
 
 
-_import_plugins()
+_import_plugins()

Diff for: pandas_style_guide/_plugins/abc.py

@@ -1,9 +1,19 @@
 import ast
-from pandas_style_guide._data import register
-from pandas_style_guide._ast_helpers import is_name_attr
+from typing import Iterator, Tuple
+
+from pandas_style_guide._data import State, register
+
+MSG = "PSG010 don't import from collections.abc"
+
 
-MSG = 'PSG010 don\'t import from collections.abc'
 @register(ast.ImportFrom)
-def abc(state, node, parent):
-    if node.module == 'collections.abc' or (node.module == 'collections' and 'abc' in {name.name for name in node.names}):
-        yield node.lineno, node.col_offset, MSG
+def visit_ImportFrom(
+    state: State,
+    node: ast.ImportFrom,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
+    if node.module == "collections.abc" or (
+        node.module == "collections"
+        and "abc" in {name.name for name in node.names}
+    ):
+        yield node.lineno, node.col_offset, MSG

Diff for: pandas_style_guide/_plugins/api_types.py

@@ -1,8 +1,23 @@
 import ast
-from pandas_style_guide._data import register
+from typing import Iterator, Tuple
+
+from pandas_style_guide._data import State, register
+
+MSG = "PSG005 dont use pd.api.types, import from pandas.api.types instead"
+
 
-MSG = 'PSG005 dont use pd.api.types, import from pandas.api.types instead'
 @register(ast.Attribute)
-def check_for_pytest_warns(state, node, parent):
-    if isinstance(node.value, ast.Attribute) and node.value.attr == 'types' and isinstance(node.value.value, ast.Attribute) and node.value.value.attr == 'api' and isinstance(node.value.value.value, ast.Name) and node.value.value.value.id in {'pd', 'pandas'}:
+def check_for_pytest_warns(
+    state: State,
+    node: ast.Attribute,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
+    if (
+        isinstance(node.value, ast.Attribute)
+        and node.value.attr == "types"
+        and isinstance(node.value.value, ast.Attribute)
+        and node.value.value.attr == "api"
+        and isinstance(node.value.value.value, ast.Name)
+        and node.value.value.value.id in {"pd", "pandas"}
+    ):
         yield node.lineno, node.col_offset, MSG

Diff for: pandas_style_guide/_plugins/bare_pytest_raises.py

@@ -1,11 +1,18 @@
 import ast
-from pandas_style_guide._data import register
-from pandas_style_guide._ast_helpers import is_name_attr
+from typing import Iterator, Tuple
+
+from pandas_style_guide._data import State, register
+
+MSG = "PSG010 bare pytest raises found"
+
 
-MSG = 'PSG010 bare pytest raises found'
 @register(ast.Call)
-def np_bool_object(state, node, parent):
+def visit_Call(
+    state: State,
+    node: ast.Call,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
     if not node.keywords:
         yield node.lineno, node.col_offset, MSG
-    elif 'match' not in {keyword.arg for keyword in node.keywords}:
+    elif "match" not in {keyword.arg for keyword in node.keywords}:
         yield node.lineno, node.col_offset, MSG

Diff for: pandas_style_guide/_plugins/builtin_filter.py

@@ -1,9 +1,17 @@
 import ast
-from pandas_style_guide._data import register
+from typing import Iterator, Tuple
+
+from pandas_style_guide._data import State, register
+
+MSG = "PSG004 do not use builtin filter function"
+
 
-MSG = 'PSG004 do not use builtin filter function'
 @register(ast.Call)
-def builtin_filter(state, node, parent):
+def visit_Call(
+    state: State,
+    node: ast.Call,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
     if isinstance(node.func, ast.Name):
-        if node.func.id == 'filter':
+        if node.func.id == "filter":
             yield (node.lineno, node.col_offset, MSG)

Diff for: pandas_style_guide/_plugins/check_contextless_pytest_raises.py

@@ -1,11 +1,33 @@
 import ast
-from pandas_style_guide._data import register
+from typing import Iterator, Tuple
+
 from pandas_style_guide._ast_helpers import is_name_attr
+from pandas_style_guide._data import State, register
+
+MSG = "PSG003 Do not use pytest.raises without context manager"
+
 
-MSG = 'PSG003 Do not use pytest.raises without context manager'
 @register(ast.Call)
-def check_contextless_pytest_raises(state, node, parent):
-    if is_name_attr(node.func, state.from_imports, 'pytest', ('raises', )) and not isinstance(parent, ast.withitem):
+def visit_Call(
+    state: State,
+    node: ast.Call,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
+    if (
+        is_name_attr(
+            node.func,
+            state.from_imports,
+            "pytest",
+            ("raises",),
+        )
+        and not isinstance(parent, ast.withitem)
+    ):
         yield node.lineno, node.col_offset, MSG
-    elif isinstance(node.func, ast.Attribute) and node.func.attr == 'raises' and isinstance(node.func.value, ast.Name) and node.func.value.id == 'pytest' and not isinstance(parent, ast.withitem):
+    elif (
+        isinstance(node.func, ast.Attribute)
+        and node.func.attr == "raises"
+        and isinstance(node.func.value, ast.Name)
+        and node.func.value.id == "pytest"
+        and not isinstance(parent, ast.withitem)
+    ):
         yield node.lineno, node.col_offset, MSG

Diff for: pandas_style_guide/_plugins/check_for_exec.py

@@ -1,9 +1,17 @@
 import ast
-from pandas_style_guide._data import register
+from typing import Iterator, Tuple
+
+from pandas_style_guide._data import State, register
+
+MSG = "PSG001 do not use exec"
+
 
-MSG = 'PSG001 do not use exec'
 @register(ast.Call)
-def check_for_exec(state, node, parent):
+def visit_Call(
+    state: State,
+    node: ast.Call,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
     if isinstance(node.func, ast.Name):
-        if node.func.id == 'exec':
+        if node.func.id == "exec":
             yield (node.lineno, node.col_offset, MSG)

Diff for: pandas_style_guide/_plugins/check_for_pytest_warns.py

@@ -1,8 +1,20 @@
 import ast
-from pandas_style_guide._data import register
+from typing import Iterator, Tuple
+
+from pandas_style_guide._data import State, register
+
+MSG = "PSG002 do not use pytest.warns"
+
 
-MSG = 'PSG002 do not use pytest.warns'
 @register(ast.Attribute)
-def check_for_pytest_warns(state, node, parent):
-    if node.attr == 'warns' and isinstance(node.value, ast.Name) and node.value.id == 'pytest' :
-        yield node.lineno, node.col_offset, MSG
+def check_for_pytest_warns(
+    state: State,
+    node: ast.Attribute,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
+    if (
+        node.attr == "warns"
+        and isinstance(node.value, ast.Name)
+        and node.value.id == "pytest"
+    ):
+        yield node.lineno, node.col_offset, MSG

Diff for: pandas_style_guide/_plugins/class.py

@@ -1,9 +1,16 @@
 import ast
-from pandas_style_guide._data import register
-from pandas_style_guide._ast_helpers import is_name_attr
+from typing import Iterator, Tuple
+
+from pandas_style_guide._data import State, register
+
+MSG = "PSG008 don't use .__class__, use type()"
+
 
-MSG = 'PSG008 don\'t use .__class__, use type()'
 @register(ast.Attribute)
-def dunder_class(state, node, parent):
-    if node.attr == '__class__':
-        yield node.lineno, node.col_offset, MSG
+def visit_Attribute(
+    state: State,
+    node: ast.Attribute,
+    parent: ast.AST,
+) -> Iterator[Tuple[int, int, str]]:
+    if node.attr == "__class__":
+        yield node.lineno, node.col_offset, MSG