check early returns + match statement

Author: rmorshea

README.md

@@ -42,6 +42,10 @@ tox
         <td>ROH102</td>
         <td>Hook was used inside a conditional or loop statement</td>
     </tr>
+    <tr>
+        <td>ROH103</td>
+        <td>Hook was used after an early return</td>
+    </tr>
     <tr>
         <td>ROH200</td>
         <td>

flake8_idom_hooks/rules_of_hooks.py

@@ -1,18 +1,21 @@
+from __future__ import annotations
+
 import ast
-from typing import Optional, Union
+import sys
 
 from .common import CheckContext, set_current
 
 
 class RulesOfHooksVisitor(ast.NodeVisitor):
     def __init__(self, context: CheckContext) -> None:
         self._context = context
-        self._current_hook: Optional[ast.FunctionDef] = None
-        self._current_component: Optional[ast.FunctionDef] = None
-        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
+        self._current_call: ast.Call | None = None
+        self._current_component: ast.FunctionDef | None = None
+        self._current_conditional: ast.If | ast.IfExp | ast.Try | None = None
+        self._current_early_return: ast.Return | None = None
+        self._current_function: ast.FunctionDef | None = None
+        self._current_hook: ast.FunctionDef | None = None
+        self._current_loop: ast.For | ast.While | None = None
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
         if self._context.is_hook_def(node):
@@ -24,6 +27,7 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
                 # we need to reset these before enter new hook
                 conditional=None,
                 loop=None,
+                early_return=None,
             ):
                 self.generic_visit(node)
         elif self._context.is_component_def(node):
@@ -34,13 +38,14 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
                 # we need to reset these before visiting a new component
                 conditional=None,
                 loop=None,
+                early_return=None,
             ):
                 self.generic_visit(node)
         else:
             with set_current(self, function=node):
                 self.generic_visit(node)
 
-    def _visit_hook_usage(self, node: Union[ast.Name, ast.Attribute]) -> None:
+    def _visit_hook_usage(self, node: ast.AST) -> None:
         self._check_if_propper_hook_usage(node)
 
     visit_Attribute = _visit_hook_usage
@@ -53,6 +58,7 @@ def _visit_conditional(self, node: ast.AST) -> None:
     visit_If = _visit_conditional
     visit_IfExp = _visit_conditional
     visit_Try = _visit_conditional
+    visit_Match = _visit_conditional
 
     def _visit_loop(self, node: ast.AST) -> None:
         with set_current(self, loop=node):
@@ -61,14 +67,15 @@ def _visit_loop(self, node: ast.AST) -> None:
     visit_For = _visit_loop
     visit_While = _visit_loop
 
+    def visit_Return(self, node: ast.Return) -> None:
+        self._current_early_return = node
+
     def _check_if_hook_defined_in_function(self, node: ast.FunctionDef) -> None:
         if self._current_function is not None:
             msg = f"hook {node.name!r} defined as closure in function {self._current_function.name!r}"
             self._context.add_error(100, node, msg)
 
-    def _check_if_propper_hook_usage(
-        self, node: Union[ast.Name, ast.Attribute]
-    ) -> None:
+    def _check_if_propper_hook_usage(self, node: ast.Name | ast.Attribute) -> None:
         if isinstance(node, ast.Name):
             name = node.id
         else:
@@ -83,14 +90,24 @@ def _check_if_propper_hook_usage(
 
         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]
+            node_name = _NODE_TYPE_TO_NAME[type(loop_or_conditional)]
             msg = f"hook {name!r} used inside {node_name}"
             self._context.add_error(102, node, msg)
+
+        if self._current_early_return:
+            self._context.add_error(
+                103,
+                node,
+                f"hook {name!r} used after an early return",
+            )
+
+
+_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",
+}
+if sys.version_info >= (3, 10):
+    _NODE_TYPE_TO_NAME[ast.Match] = "match statement"

tests/cases/hook_usage.py

@@ -159,3 +159,11 @@ def Component():
     @component
     def use_other():
         use_state
+
+
+@component
+def example():
+    if True:
+        return None
+    # error: ROH103 hook 'use_state' used after an early return
+    use_state()

tests/cases/match_statement.py

@@ -0,0 +1,5 @@
+@component
+def example():
+    match something:
+        case int:
+            use_state()

tests/test_cases.py

@@ -1,5 +1,6 @@
 import ast
 from pathlib import Path
+import sys
 
 import pytest
 from flake8.options.manager import OptionManager
@@ -32,14 +33,22 @@ def setup_plugin(args):
             "",
             "hook_usage.py",
         ),
-        (
-            "--exhaustive-hook-deps",
-            "exhaustive_deps.py",
-        ),
         (
             "",
             "no_exhaustive_deps.py",
         ),
+        pytest.param(
+            "",
+            "match_statement.py",
+            marks=pytest.mark.skipif(
+                sys.version_info < (3, 10),
+                reason="Match statement only in Python 3.10 and above",
+            ),
+        ),
+        (
+            "--exhaustive-hook-deps",
+            "exhaustive_deps.py",
+        ),
         (
             r"--component-decorator-pattern ^(component|custom_component)$",
             "custom_component_decorator_pattern.py",