fix bug where nested element def treated as hook

Author: rmorshea

flake8_idom_hooks/exhaustive_deps.py

@@ -1,7 +1,7 @@
 import ast
 from typing import Optional, List, Union, Set
 
-from .utils import is_hook_or_element_def, ErrorVisitor
+from .utils import is_hook_def, is_element_def, ErrorVisitor
 
 
 HOOKS_WITH_DEPS = ("use_effect", "use_callback", "use_memo")
@@ -13,7 +13,7 @@ def __init__(self) -> None:
         self._current_hook_or_element: Optional[ast.FunctionDef] = None
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
-        if is_hook_or_element_def(node):
+        if is_hook_def(node) or is_element_def(node):
             self._current_hook_or_element = node
             self.generic_visit(node)
             self._current_hook_or_element = None

flake8_idom_hooks/rules_of_hooks.py

@@ -2,22 +2,26 @@
 from contextlib import contextmanager
 from typing import Iterator, Union, Optional, Any
 
-from .utils import is_hook_or_element_def, ErrorVisitor, is_hook_function_name
+from .utils import is_hook_def, is_element_def, ErrorVisitor, is_hook_function_name
 
 
 class RulesOfHooksVisitor(ErrorVisitor):
     def __init__(self) -> None:
         super().__init__()
-        self._current_hook_or_element: Optional[ast.FunctionDef] = None
+        self._current_hook: Optional[ast.FunctionDef] = None
+        self._current_element: 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
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
-        if is_hook_or_element_def(node):
+        if is_hook_def(node):
             self._check_if_hook_defined_in_function(node)
-            with self._set_current(hook_or_element=node, function=node):
+            with self._set_current(hook=node, function=node):
+                self.generic_visit(node)
+        elif is_element_def(node):
+            with self._set_current(element=node, function=node):
                 self.generic_visit(node)
         else:
             with self._set_current(function=node):
@@ -60,7 +64,7 @@ def _check_if_propper_hook_usage(
         if not is_hook_function_name(name):
             return None
 
-        if self._current_hook_or_element is None:
+        if self._current_hook is None and self._current_element is None:
             msg = f"hook {name!r} used outside element or hook definition"
             self._save_error(101, node, msg)
 

flake8_idom_hooks/utils.py

@@ -10,8 +10,12 @@ def _save_error(self, error_code: int, node: ast.AST, message: str) -> None:
         self.errors.append((node.lineno, node.col_offset, f"ROH{error_code} {message}"))
 
 
-def is_hook_or_element_def(node: ast.FunctionDef) -> bool:
-    return is_element_function_name(node.name) or is_hook_function_name(node.name)
+def is_hook_def(node: ast.FunctionDef) -> bool:
+    return is_hook_function_name(node.name)
+
+
+def is_element_def(node: ast.FunctionDef) -> bool:
+    return is_element_function_name(node.name)
 
 
 def is_element_function_name(name: str) -> bool:

tests/hook_usage_test_cases.py

@@ -200,3 +200,9 @@ def impropper_usage_of_effect_as_decorator():
         # error: ROH202 dependency args of 'use_effect' should be a literal list or tuple - not expression type 'Name'
         not_a_list_or_tuple,
     )
+
+
+def make_element():
+    # nested element definitions are ok.
+    def NestedElement():
+        use_state