initial implementation and testing

Author: rmorshea

flake8_idom_hooks.py

@@ -1,7 +1,10 @@
+import ast
+from contextlib import contextmanager
 from pkg_resources import (
     get_distribution as _get_distribution,
     DistributionNotFound as _DistributionNotFound,
 )
+from typing import List, Tuple, Iterator, Union, Optional
 
 
 try:
@@ -17,8 +20,106 @@ class Plugin:
     version = __version__
     options = None
 
-    def __init__(self, tree):
+    def __init__(self, tree: ast.Module):
         self._tree = tree
 
     def run(self):
-        return [(7, 39, "hooks cannot be used in conditionals", type(self))]
+        visitor = HookRulesVisitor()
+        visitor.visit(self._tree)
+        cls = type(self)
+        return [(line, col, msg, cls) for line, col, msg in visitor.errors]
+
+
+class HookRulesVisitor(ast.NodeVisitor):
+    def __init__(self):
+        self.errors: List[Tuple[int, int, str]] = []
+        self._current_function: Optional[ast.FunctionDef] = None
+        self._current_call: Optional[ast.Call] = None
+        self._current_conditional: Union[None, ast.If, ast.IfExp, ast.Try] = None
+        self._current_loop: Union[None, ast.For, ast.While] = None
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        self._check_if_hook_defined_in_function(node)
+        with self._set_current(function=node):
+            self.generic_visit(node)
+
+    def _visit_hook_usage(self, node: ast.AST) -> None:
+        self._check_if_propper_hook_usage(node)
+
+    visit_Attribute = _visit_hook_usage
+    visit_Name = _visit_hook_usage
+
+    def _visit_conditional(self, node: ast.AST) -> None:
+        with self._set_current(conditional=node):
+            self.generic_visit(node)
+
+    visit_If = _visit_conditional
+    visit_IfExp = _visit_conditional
+    visit_Try = _visit_conditional
+
+    def _visit_loop(self, node: ast.AST) -> None:
+        with self._set_current(loop=node):
+            self.generic_visit(node)
+
+    visit_For = _visit_loop
+    visit_While = _visit_loop
+
+    def _check_if_hook_defined_in_function(self, node: ast.FunctionDef) -> None:
+        if self._current_function is not None and _is_hook_or_element_def(node):
+            msg = f"Hook {node.name!r} defined inside another function."
+            self.errors.append((node.lineno, node.col_offset, msg))
+
+    def _check_if_propper_hook_usage(self, node: Union[ast.Name, ast.Attribute]):
+        if isinstance(node, ast.Name):
+            name = node.id
+        else:
+            name = node.attr
+
+        if not _is_hook_function_name(name):
+            return
+
+        if not _is_hook_or_element_def(self._current_function):
+            msg = f"Hook {name!r} used outside element or hook definition."
+            self.errors.append((node.lineno, node.col_offset, msg))
+            return
+
+        _loop_or_conditional = self._current_conditional or self._current_loop
+        if _loop_or_conditional is not None:
+            node_type = type(_loop_or_conditional)
+            node_type_to_name = {
+                ast.If: "if statement",
+                ast.IfExp: "inline if expression",
+                ast.Try: "try statement",
+                ast.For: "for loop",
+                ast.While: "while loop",
+            }
+            node_name = node_type_to_name[node_type]
+            msg = f"Hook {name!r} used inside {node_name}."
+            self.errors.append((node.lineno, node.col_offset, msg))
+            return
+
+    @contextmanager
+    def _set_current(self, **attrs) -> Iterator[None]:
+        old_attrs = {k: getattr(self, f"_current_{k}") for k in attrs}
+        for k, v in attrs.items():
+            setattr(self, f"_current_{k}", v)
+        try:
+            yield
+        finally:
+            for k, v in old_attrs.items():
+                setattr(self, f"_current_{k}", v)
+
+
+def _is_hook_or_element_def(node: Optional[ast.FunctionDef]) -> bool:
+    if node is None:
+        return False
+    else:
+        return _is_element_function_name(node.name) or _is_hook_function_name(node.name)
+
+
+def _is_element_function_name(name: str) -> bool:
+    return name[0].upper() == name[0] and "_" not in name
+
+
+def _is_hook_function_name(name: str) -> bool:
+    return name.lstrip("_").startswith("use_")

setup.cfg

@@ -6,6 +6,7 @@ select = B,C,E,F,W,T4,B9,N
 exclude =
     .eggs/*
     .tox/*
+    tests/hook_usage_cases.py
 
 [coverage:report]
 fail_under = 100

tests/cases/hook_in_conditional.py

@@ -0,0 +1,101 @@
+def HookInIf():
+    if True:
+        # error: Hook 'use_state' used inside if statement.
+        use_state
+
+
+def HookInElif():
+    if False:
+        pass
+    elif True:
+        # error: Hook 'use_state' used inside if statement.
+        use_state
+
+
+def HookInElse():
+    if False:
+        pass
+    else:
+        # error: Hook 'use_state' used inside if statement.
+        use_state
+
+
+def HookInIfExp():
+    (
+        # error: Hook 'use_state' used inside inline if expression.
+        use_state
+        if True
+        else None
+    )
+
+
+def HookInElseOfIfExp():
+    (
+        None
+        if True
+        else
+        # error: Hook 'use_state' used inside inline if expression.
+        use_state
+    )
+
+
+def HookInTry():
+    try:
+        # error: Hook 'use_state' used inside try statement.
+        use_state
+    except:
+        pass
+
+
+def HookInExcept():
+    try:
+        raise ValueError()
+    except:
+        # error: Hook 'use_state' used inside try statement.
+        use_state
+
+
+def HookInFinally():
+    try:
+        pass
+    finally:
+        # error: Hook 'use_state' used inside try statement.
+        use_state
+
+
+def HookInForLoop():
+    for i in range(3):
+        # error: Hook 'use_state' used inside for loop.
+        use_state
+
+
+def HookInWhileLoop():
+    while True:
+        # error: Hook 'use_state' used inside while loop.
+        use_state
+
+
+def outer_function():
+    # error: Hook 'use_state' defined inside another function.
+    def use_state():
+        ...
+
+
+def generic_function():
+    # error: Hook 'use_state' used outside element or hook definition.
+    use_state
+
+
+def use_state():
+    use_other
+
+
+def Element():
+    use_state
+
+
+# ok since 'use_state' is not the last attribute
+module.use_state.other
+
+# error: Hook 'use_state' used outside element or hook definition.
+module.use_state

tests/hook_usage_cases.py

@@ -1,27 +1,24 @@
 import ast
 from pathlib import Path
 
-import pytest
-
 from flake8_idom_hooks import Plugin
 
 
-case_file_asts = {}
-for path in (Path(__file__).parent / "cases").rglob("*.py"):
-    with path.open() as f:
-        tree = ast.parse(f.read(), path.name)
-        case_file_asts[path.name[:-3]] = tree
-
-
-expectations = {
-    "hook_in_conditional": [(7, 39, "hooks cannot be used in conditionals")]
-}
-
-
-@pytest.mark.parametrize(
-    "tree, expected_errors",
-    [(case_file_asts[name], expectations[name]) for name in case_file_asts],
-)
-def test_flake8_idom_hooks(tree, expected_errors):
+def test_flake8_idom_hooks():
+    path_to_case_file = Path(__file__).parent / "hook_usage_cases.py"
+    with path_to_case_file.open() as file:
+        # save the file's AST
+        file_content = file.read()
+        tree = ast.parse(file_content, path_to_case_file.name)
+
+        # find 'error' comments to construct expectations
+        expected_errors = []
+        for index, line in enumerate(file_content.split("\n")):
+            lstrip_line = line.lstrip()
+            if lstrip_line.startswith("# error:"):
+                lineno = index + 2  # use 2 since error should be on next line
+                col_offset = len(line) - len(lstrip_line)
+                message = line.replace("# error:", "", 1).strip()
+                expected_errors.append((lineno, col_offset, message))
     actual_errors = Plugin(tree).run()
     assert [(ln, col, msg) for ln, col, msg, p_type in actual_errors] == expected_errors