Add filter function to list

run-ai · Dec 19, 2024 · 65c5e61 · 65c5e61
1 parent 48eb474
commit 65c5e61
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/py/runai_model_streamer/runai_model_streamer/file_streamer/file_streamer.py b/py/runai_model_streamer/runai_model_streamer/file_streamer/file_streamer.py
@@ -102,9 +102,9 @@ def request_ready_chunks(self) -> Iterator:
                 self.current_request_chunks[:relative_index]
             )
 
-    def list(self, path: str, pattern: str) -> List[str]:
+    def list(self, path: str, pattern: str = "*.*") -> List[str]:
         all_files = runai_list(self.streamer, path)
         return [file for file in all_files if fnmatch.fnmatch(file, pattern)]
 
-    def copy(self, src_path : str, dst_path : str) -> None:
-        return runai_copy(self.streamer, src_path, dst_path)
+    def copy(self, src_path: str, dst_path: str) -> None:
+        return runai_copy(self.streamer, src_path, dst_path)
diff --git a/py/runai_model_streamer/runai_model_streamer/file_streamer/tests/test_file_streamer.py b/py/runai_model_streamer/runai_model_streamer/file_streamer/tests/test_file_streamer.py
@@ -99,6 +99,13 @@ def test_limited_memory_cap(self, mock_get_memory_mode, mock_getenv):
                     id_to_results[id]["expected_text"],
                 )
 
+    @patch("runai_model_streamer.file_streamer.file_streamer.runai_list")
+    def test_list_pattern(self, mock_runai_list):
+        mock_runai_list.return_value = ["a.safetensors", "b.py", "c.safetensors"]
+        with FileStreamer() as fs:
+            filtered = fs.list("test", "*.safetensors")
+            self.assertCountEqual(filtered, ["c.safetensors", "a.safetensors"])
+
     def tearDown(self):
         shutil.rmtree(self.temp_dir)