diff --git a/codeflash/code_utils/static_analysis.py b/codeflash/code_utils/static_analysis.py index ad0fa3459..7abbcefcd 100644 --- a/codeflash/code_utils/static_analysis.py +++ b/codeflash/code_utils/static_analysis.py @@ -7,6 +7,8 @@ from pydantic import BaseModel, ConfigDict, field_validator +from codeflash.models.models import FunctionParent + if TYPE_CHECKING: from codeflash.models.models import FunctionParent @@ -139,16 +141,19 @@ def get_first_top_level_function_or_method_ast( def function_kind(node: ast.FunctionDef | ast.AsyncFunctionDef, parents: list[FunctionParent]) -> FunctionKind | None: - if not parents or parents[0].type in ["FunctionDef", "AsyncFunctionDef"]: + if not parents: + return FunctionKind.FUNCTION + parent_type = parents[0].type + if parent_type in {"FunctionDef", "AsyncFunctionDef"}: return FunctionKind.FUNCTION - for _i in range(len(parents) - 1, -1, -1): - continue - if parents[0].type == "ClassDef": + if parent_type == "ClassDef": + # Fast path: decorator_list is typically small; scan for 'classmethod' and 'staticmethod' for decorator in node.decorator_list: if isinstance(decorator, ast.Name): - if decorator.id == "classmethod": + did = decorator.id + if did == "classmethod": return FunctionKind.CLASS_METHOD - if decorator.id == "staticmethod": + if did == "staticmethod": return FunctionKind.STATIC_METHOD return FunctionKind.INSTANCE_METHOD return None