From 99189a7cb20093be2046c50619cb06db98615690 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 1 Apr 2025 15:26:40 -0700 Subject: [PATCH] prevent nested classes/ functions from showing up in code context . helper_functions --- codeflash/context/code_context_extractor.py | 4 +++- tests/test_code_context_extractor.py | 10 ++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/codeflash/context/code_context_extractor.py b/codeflash/context/code_context_extractor.py index 3827239ed..13c7ec1dd 100644 --- a/codeflash/context/code_context_extractor.py +++ b/codeflash/context/code_context_extractor.py @@ -410,10 +410,12 @@ def get_function_sources_from_jedi( and definition.full_name and definition.type == "function" and not belongs_to_function_qualified(definition, qualified_function_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 ): function_source = FunctionSource( file_path=definition_path, - qualified_name=get_qualified_name(definition.module_name, definition.full_name), + qualified_name=qualified_name, fully_qualified_name=definition.full_name, only_function_name=definition.name, source_code=definition.get_line_code(), diff --git a/tests/test_code_context_extractor.py b/tests/test_code_context_extractor.py index 65c636582..069a8eb19 100644 --- a/tests/test_code_context_extractor.py +++ b/tests/test_code_context_extractor.py @@ -23,6 +23,12 @@ def innocent_bystander(self): def helper_method(self): return self.name + class NestedClass: + def __init__(self, name): + self.name = name + + def nested_method(self): + return self.name def main_method(): return "hello" @@ -33,6 +39,7 @@ def __init__(self, name): self.name = name def main_method(self): + self.name = HelperClass.NestedClass("test").nested_method() return HelperClass(self.name).helper_method() @@ -73,6 +80,8 @@ def test_code_replacement10() -> None: ) code_ctx = get_code_optimization_context(function_to_optimize=func_top_optimize, project_root_path=file_path.parent) + qualified_names = {func.qualified_name for func in code_ctx.helper_functions} + assert qualified_names == {"HelperClass.helper_method"} # Nested method should not be in here read_write_context, read_only_context = code_ctx.read_writable_code, code_ctx.read_only_context_code expected_read_write_context = """ @@ -91,6 +100,7 @@ def __init__(self, name): self.name = name def main_method(self): + self.name = HelperClass.NestedClass("test").nested_method() return HelperClass(self.name).helper_method() """ expected_read_only_context = """