From 2c758ade861d50d00190e65b8839f816665698de Mon Sep 17 00:00:00 2001 From: Edward Li Date: Fri, 14 Mar 2025 15:43:07 -0700 Subject: [PATCH 1/2] Speed up list_directory tool --- src/codegen/extensions/tools/list_directory.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/codegen/extensions/tools/list_directory.py b/src/codegen/extensions/tools/list_directory.py index b2668c435..243c79c87 100644 --- a/src/codegen/extensions/tools/list_directory.py +++ b/src/codegen/extensions/tools/list_directory.py @@ -140,9 +140,8 @@ def get_directory_info(dir_obj: Directory, current_depth: int) -> DirectoryInfo: """Helper function to get directory info recursively.""" # Get direct files (always include files unless at max depth) all_files = [] - for file in dir_obj.files(recursive=True): - if file.directory == dir_obj: - all_files.append(file.filepath.split("/")[-1]) + for file_name in dir_obj._files: + all_files.append(file_name) # Get direct subdirectories subdirs = [] From a40ee36bc2f2d85e17738e08630de472ba755791 Mon Sep 17 00:00:00 2001 From: Edward Li Date: Fri, 14 Mar 2025 15:47:16 -0700 Subject: [PATCH 2/2] clean up implementation --- src/codegen/extensions/tools/list_directory.py | 2 +- src/codegen/sdk/core/directory.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/codegen/extensions/tools/list_directory.py b/src/codegen/extensions/tools/list_directory.py index 243c79c87..357f303ca 100644 --- a/src/codegen/extensions/tools/list_directory.py +++ b/src/codegen/extensions/tools/list_directory.py @@ -140,7 +140,7 @@ def get_directory_info(dir_obj: Directory, current_depth: int) -> DirectoryInfo: """Helper function to get directory info recursively.""" # Get direct files (always include files unless at max depth) all_files = [] - for file_name in dir_obj._files: + for file_name in dir_obj.file_names: all_files.append(file_name) # Get direct subdirectories diff --git a/src/codegen/sdk/core/directory.py b/src/codegen/sdk/core/directory.py index 504d608e0..806e90ff8 100644 --- a/src/codegen/sdk/core/directory.py +++ b/src/codegen/sdk/core/directory.py @@ -196,6 +196,11 @@ def item_names(self, recursive: bool = False) -> list[str]: """ return self._files + self._subdirectories + @property + def file_names(self) -> list[str]: + """Get a list of all file names in the directory.""" + return self._files + @property def tree(self) -> list[Self | TFile]: """Get a recursive list of all files and subdirectories in the directory.