type checks

Author: KRRT7

codeflash/code_utils/code_extractor.py

@@ -655,8 +655,10 @@ def find_target(node_list: list[ast.stmt], name_parts: tuple[str, str] | tuple[s
         if target is None or len(name_parts) == 1:
             return target
 
-        if not isinstance(target, ast.ClassDef):
+        if not isinstance(target, ast.ClassDef) or len(name_parts) < 2:
             return None
+        # At this point, name_parts has at least 2 elements
+        method_name: str = name_parts[1]  # type: ignore[misc]
         class_skeleton.add((target.lineno, target.body[0].lineno - 1))
         cbody = target.body
         if isinstance(cbody[0], ast.expr):  # Is a docstring
@@ -669,15 +671,15 @@ def find_target(node_list: list[ast.stmt], name_parts: tuple[str, str] | tuple[s
             if (
                 isinstance(cnode, (ast.FunctionDef, ast.AsyncFunctionDef))
                 and len(cnode_name := cnode.name) > 4
-                and cnode_name != name_parts[1]
+                and cnode_name != method_name
                 and cnode_name.isascii()
                 and cnode_name.startswith("__")
                 and cnode_name.endswith("__")
             ):
                 contextual_dunder_methods.add((target.name, cnode_name))
                 class_skeleton.add((cnode.lineno, cnode.end_lineno))
 
-        return find_target(target.body, name_parts[1:])
+        return find_target(target.body, (method_name,))
 
     with file_path.open(encoding="utf8") as file:
         source_code: str = file.read()
@@ -708,9 +710,14 @@ def find_target(node_list: list[ast.stmt], name_parts: tuple[str, str] | tuple[s
         )
         return None, set()
     for qualified_name_parts in qualified_name_parts_list:
-        target_node: ast.AST | None = find_target(module_node.body, qualified_name_parts)
+        target_node = find_target(module_node.body, qualified_name_parts)
         if target_node is None:
             continue
+        # find_target returns FunctionDef, AsyncFunctionDef, ClassDef, Assign, or AnnAssign - all have lineno/end_lineno
+        if not isinstance(
+            target_node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Assign, ast.AnnAssign)
+        ):
+            continue
 
         if (
             isinstance(target_node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))