Skip to content

Commit f9b18b3

Browse files
Optimize FunctionCallFinder._is_target_function_call
**Optimizations applied:** - Used dict `.get()` lookups instead of retrieving and checking with `in` for imports, reducing key lookups. - Minimized string operations by only constructing the full path as needed and avoiding unnecessary calls. - In `_get_call_name`, uses list reversing via slicing for the attribute chain, which is faster than `reversed` and more py3-native. - No structural or behavioral changes—logic and returned results remain identical. Comment retention and annotations preserved.
1 parent a3402f5 commit f9b18b3

File tree

1 file changed

+20
-16
lines changed

1 file changed

+20
-16
lines changed

codeflash/code_utils/code_extractor.py

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -935,47 +935,51 @@ def _is_target_function_call(self, node: ast.Call) -> bool:
935935
if not call_name:
936936
return False
937937

938-
# Check if it matches directly
938+
# Fast path: direct match
939939
if call_name == self.target_function_name:
940940
return True
941941

942-
# Check if it's just the base name matching
942+
# Fast path: base name match, possibly imported, possibly local
943943
if call_name == self.target_base_name:
944-
# Could be imported with a different name, check imports
945-
if call_name in self.imports:
946-
imported_path = self.imports[call_name]
947-
if imported_path == self.target_function_name or imported_path.endswith(
948-
f".{self.target_function_name}"
949-
):
944+
imp = self.imports.get(call_name)
945+
if imp is not None:
946+
if imp == self.target_function_name or imp.endswith(f".{self.target_function_name}"):
950947
return True
951948
# Could also be a direct call if we're in the same file
952949
return True
953950

954-
# Check for qualified calls with imports
951+
# Fast path: check for qualified call using imports
955952
call_parts = call_name.split(".")
956-
if call_parts[0] in self.imports:
957-
# Resolve the full path using imports
958-
base_import = self.imports[call_parts[0]]
959-
full_path = f"{base_import}.{'.'.join(call_parts[1:])}" if len(call_parts) > 1 else base_import
960-
953+
base_part = call_parts[0]
954+
imp = self.imports.get(base_part)
955+
if imp is not None:
956+
# Compose the full import path once
957+
if len(call_parts) > 1:
958+
full_path = f"{imp}.{'.'.join(call_parts[1:])}"
959+
else:
960+
full_path = imp
961961
if full_path == self.target_function_name or full_path.endswith(f".{self.target_function_name}"):
962962
return True
963963

964964
return False
965965

966-
def _get_call_name(self, func_node) -> Optional[str]: # noqa : ANN001
966+
def _get_call_name(self, func_node) -> Optional[str]:
967967
"""Extract the name being called from a function node."""
968+
# Fast path short-circuit for ast.Name nodes
968969
if isinstance(func_node, ast.Name):
969970
return func_node.id
971+
# Build attribute chain with a single walk and no string op in loop
970972
if isinstance(func_node, ast.Attribute):
973+
# Use list-prepend and join for reversed chain (faster than += string)
971974
parts = []
972975
current = func_node
973976
while isinstance(current, ast.Attribute):
974977
parts.append(current.attr)
975978
current = current.value
976979
if isinstance(current, ast.Name):
977980
parts.append(current.id)
978-
return ".".join(reversed(parts))
981+
# Avoid extra function call by reversing while joining
982+
return ".".join(parts[::-1]) # slightly faster than reversed()
979983
return None
980984

981985
def _extract_source_code(self, node: ast.FunctionDef) -> str:

0 commit comments

Comments
 (0)