Add missing type hints.

comfyanonymous · Aug 13, 2024 · e7b501f · e7b501f
1 parent b3b7b9a
commit e7b501f
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/api_server/services/file_service.py b/api_server/services/file_service.py
@@ -1,12 +1,13 @@
-from api_server.utils.file_operations import FileSystemOperations
+from typing import Dict, List, Optional
+from api_server.utils.file_operations import FileSystemOperations, FileSystemItem
 
 class FileService:
-    def __init__(self, allowed_directories, file_system_ops=None):
-        self.allowed_directories = allowed_directories
-        self.file_system_ops = file_system_ops or FileSystemOperations()
+    def __init__(self, allowed_directories: Dict[str, str], file_system_ops: Optional[FileSystemOperations] = None):
+        self.allowed_directories: Dict[str, str] = allowed_directories
+        self.file_system_ops: FileSystemOperations = file_system_ops or FileSystemOperations()
 
-    def list_files(self, directory_key):
+    def list_files(self, directory_key: str) -> List[FileSystemItem]:
         if directory_key not in self.allowed_directories:
             raise ValueError("Invalid directory key")
-        directory_path = self.allowed_directories[directory_key]
+        directory_path: str = self.allowed_directories[directory_key]
         return self.file_system_ops.walk_directory(directory_path)