diff --git a/codeflash/context/code_context_extractor.py b/codeflash/context/code_context_extractor.py index 13c7ec1dd..516f3c94e 100644 --- a/codeflash/context/code_context_extractor.py +++ b/codeflash/context/code_context_extractor.py @@ -356,6 +356,7 @@ def get_function_to_optimize_as_function_source( name.type == "function" and name.full_name and name.name == function_to_optimize.function_name + and name.full_name.startswith(name.module_name) and get_qualified_name(name.module_name, name.full_name) == function_to_optimize.qualified_name ): function_source = FunctionSource( @@ -410,6 +411,7 @@ def get_function_sources_from_jedi( and definition.full_name and definition.type == "function" and not belongs_to_function_qualified(definition, qualified_function_name) + and definition.full_name.startswith(definition.module_name) # Avoid nested functions or classes. Only class.function is allowed and len((qualified_name := get_qualified_name(definition.module_name, definition.full_name)).split(".")) <= 2 ): diff --git a/codeflash/optimization/function_context.py b/codeflash/optimization/function_context.py index 3c28a92db..d55aa2dec 100644 --- a/codeflash/optimization/function_context.py +++ b/codeflash/optimization/function_context.py @@ -31,7 +31,7 @@ def belongs_to_class(name: Name, class_name: str) -> bool: def belongs_to_function_qualified(name: Name, qualified_function_name: str) -> bool: """Check if the given jedi Name is a direct child of the specified function, matched by qualified function name.""" try: - if get_qualified_name(name.module_name, name.full_name) == qualified_function_name: + if name.full_name.startswith(name.module_name) and get_qualified_name(name.module_name, name.full_name) == qualified_function_name: # Handles function definition and recursive function calls return False if name := name.parent():