|
7 | 7 |
|
8 | 8 | from pydantic import BaseModel, ConfigDict, field_validator |
9 | 9 |
|
| 10 | +from codeflash.models.models import FunctionParent |
| 11 | + |
10 | 12 | if TYPE_CHECKING: |
11 | 13 | from codeflash.models.models import FunctionParent |
12 | 14 |
|
@@ -139,14 +141,20 @@ def get_first_top_level_function_or_method_ast( |
139 | 141 |
|
140 | 142 |
|
141 | 143 | def function_kind(node: ast.FunctionDef | ast.AsyncFunctionDef, parents: list[FunctionParent]) -> FunctionKind | None: |
142 | | - if not parents or parents[0].type in ["FunctionDef", "AsyncFunctionDef"]: |
| 144 | + # cache parents[0].type to avoid repeated lookup |
| 145 | + if not parents: |
| 146 | + return FunctionKind.FUNCTION |
| 147 | + parent_type = parents[0].type |
| 148 | + if parent_type == "FunctionDef" or parent_type == "AsyncFunctionDef": |
143 | 149 | return FunctionKind.FUNCTION |
144 | | - if parents[0].type == "ClassDef": |
| 150 | + if parent_type == "ClassDef": |
| 151 | + # Use set for decorator comparison and break early |
145 | 152 | for decorator in node.decorator_list: |
146 | 153 | if isinstance(decorator, ast.Name): |
147 | | - if decorator.id == "classmethod": |
| 154 | + name = decorator.id |
| 155 | + if name == "classmethod": |
148 | 156 | return FunctionKind.CLASS_METHOD |
149 | | - if decorator.id == "staticmethod": |
| 157 | + if name == "staticmethod": |
150 | 158 | return FunctionKind.STATIC_METHOD |
151 | 159 | return FunctionKind.INSTANCE_METHOD |
152 | 160 | return None |
|
0 commit comments