Add GET /internal/files. (#4295)

* Create internal route table.

* List files.

* Add GET /internal/files.

Retrieves list of files in models, output, and user directories.

* Refactor file names.

* Use typing_extensions for Python 3.8

* Fix tests.

* Remove print statements.

* Update README.

* Add output and user to valid directory test.

* Add missing type hints.

api_server/routes/internal/README.md

@@ -0,0 +1,3 @@
+# ComfyUI Internal Routes
+
+All routes under the `/internal` path are designated for **internal use by ComfyUI only**. These routes are not intended for use by external applications may change at any time without notice.

api_server/routes/internal/internal_routes.py

@@ -0,0 +1,40 @@
+from aiohttp import web
+from typing import Optional
+from folder_paths import models_dir, user_directory, output_directory
+from api_server.services.file_service import FileService
+
+class InternalRoutes:
+    '''
+    The top level web router for internal routes: /internal/*
+    The endpoints here should NOT be depended upon. It is for ComfyUI frontend use only.
+    Check README.md for more information.
+    
+    '''
+    def __init__(self):
+        self.routes: web.RouteTableDef = web.RouteTableDef()
+        self._app: Optional[web.Application] = None
+        self.file_service = FileService({
+            "models": models_dir,
+            "user": user_directory,
+            "output": output_directory
+        })
+
+    def setup_routes(self):
+        @self.routes.get('/files')
+        async def list_files(request):
+            directory_key = request.query.get('directory', '')
+            try:
+                file_list = self.file_service.list_files(directory_key)
+                return web.json_response({"files": file_list})
+            except ValueError as e:
+                return web.json_response({"error": str(e)}, status=400)
+            except Exception as e:
+                return web.json_response({"error": str(e)}, status=500)
+
+
+    def get_app(self):
+        if self._app is None:
+            self._app = web.Application()
+            self.setup_routes()
+            self._app.add_routes(self.routes)
+        return self._app

api_server/services/file_service.py

@@ -0,0 +1,13 @@
+from typing import Dict, List, Optional
+from api_server.utils.file_operations import FileSystemOperations, FileSystemItem
+
+class FileService:
+    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: str) -> List[FileSystemItem]:
+        if directory_key not in self.allowed_directories:
+            raise ValueError("Invalid directory key")
+        directory_path: str = self.allowed_directories[directory_key]
+        return self.file_system_ops.walk_directory(directory_path)

api_server/utils/file_operations.py

@@ -0,0 +1,42 @@
+import os
+from typing import List, Union, TypedDict, Literal
+from typing_extensions import TypeGuard
+class FileInfo(TypedDict):
+    name: str
+    path: str
+    type: Literal["file"]
+    size: int
+
+class DirectoryInfo(TypedDict):
+    name: str
+    path: str
+    type: Literal["directory"]
+
+FileSystemItem = Union[FileInfo, DirectoryInfo]
+
+def is_file_info(item: FileSystemItem) -> TypeGuard[FileInfo]:
+    return item["type"] == "file"
+
+class FileSystemOperations:
+    @staticmethod
+    def walk_directory(directory: str) -> List[FileSystemItem]:
+        file_list: List[FileSystemItem] = []
+        for root, dirs, files in os.walk(directory):
+            for name in files:
+                file_path = os.path.join(root, name)
+                relative_path = os.path.relpath(file_path, directory)
+                file_list.append({
+                    "name": name,
+                    "path": relative_path,
+                    "type": "file",
+                    "size": os.path.getsize(file_path)
+                })
+            for name in dirs:
+                dir_path = os.path.join(root, name)
+                relative_path = os.path.relpath(dir_path, directory)
+                file_list.append({
+                    "name": name,
+                    "path": relative_path,
+                    "type": "directory"
+                })
+        return file_list

server.py

@@ -29,6 +29,8 @@ from app.frontend_management import FrontendManager
 from app.user_manager import UserManager
 from model_filemanager import download_model, DownloadModelStatus
 from typing import Optional
+from api_server.routes.internal.internal_routes import InternalRoutes
+
 
 class BinaryEventTypes:
     PREVIEW_IMAGE = 1
@@ -72,6 +74,7 @@ class PromptServer():
         mimetypes.types_map['.js'] = 'application/javascript; charset=utf-8'
 
         self.user_manager = UserManager()
+        self.internal_routes = InternalRoutes()
         self.supports = ["custom_nodes_from_web"]
         self.prompt_queue = None
         self.loop = loop
@@ -602,6 +605,7 @@ class PromptServer():
 
     def add_routes(self):
         self.user_manager.add_routes(self.routes)
+        self.app.add_subapp('/internal', self.internal_routes.get_app())
 
         # Prefix every route with /api for easier matching for delegation.
         # This is very useful for frontend dev server, which need to forward

tests-unit/server/routes/internal_routes_test.py

@@ -0,0 +1,115 @@
+import pytest
+from aiohttp import web
+from unittest.mock import MagicMock, patch
+from api_server.routes.internal.internal_routes import InternalRoutes
+from api_server.services.file_service import FileService
+from folder_paths import models_dir, user_directory, output_directory
+
+
+@pytest.fixture
+def internal_routes():
+    return InternalRoutes()
+
+@pytest.fixture
+def aiohttp_client_factory(aiohttp_client, internal_routes):
+    async def _get_client():
+        app = internal_routes.get_app()
+        return await aiohttp_client(app)
+    return _get_client
+
+@pytest.mark.asyncio
+async def test_list_files_valid_directory(aiohttp_client_factory, internal_routes):
+    mock_file_list = [
+        {"name": "file1.txt", "path": "file1.txt", "type": "file", "size": 100},
+        {"name": "dir1", "path": "dir1", "type": "directory"}
+    ]
+    internal_routes.file_service.list_files = MagicMock(return_value=mock_file_list)
+    client = await aiohttp_client_factory()
+    resp = await client.get('/files?directory=models')
+    assert resp.status == 200
+    data = await resp.json()
+    assert 'files' in data
+    assert len(data['files']) == 2
+    assert data['files'] == mock_file_list
+
+    # Check other valid directories
+    resp = await client.get('/files?directory=user')
+    assert resp.status == 200
+    resp = await client.get('/files?directory=output')
+    assert resp.status == 200
+
+@pytest.mark.asyncio
+async def test_list_files_invalid_directory(aiohttp_client_factory, internal_routes):
+    internal_routes.file_service.list_files = MagicMock(side_effect=ValueError("Invalid directory key"))
+    client = await aiohttp_client_factory()
+    resp = await client.get('/files?directory=invalid')
+    assert resp.status == 400
+    data = await resp.json()
+    assert 'error' in data
+    assert data['error'] == "Invalid directory key"
+
+@pytest.mark.asyncio
+async def test_list_files_exception(aiohttp_client_factory, internal_routes):
+    internal_routes.file_service.list_files = MagicMock(side_effect=Exception("Unexpected error"))
+    client = await aiohttp_client_factory()
+    resp = await client.get('/files?directory=models')
+    assert resp.status == 500
+    data = await resp.json()
+    assert 'error' in data
+    assert data['error'] == "Unexpected error"
+
+@pytest.mark.asyncio
+async def test_list_files_no_directory_param(aiohttp_client_factory, internal_routes):
+    mock_file_list = []
+    internal_routes.file_service.list_files = MagicMock(return_value=mock_file_list)
+    client = await aiohttp_client_factory()
+    resp = await client.get('/files')
+    assert resp.status == 200
+    data = await resp.json()
+    assert 'files' in data
+    assert len(data['files']) == 0
+
+def test_setup_routes(internal_routes):
+    internal_routes.setup_routes()
+    routes = internal_routes.routes
+    assert any(route.method == 'GET' and str(route.path) == '/files' for route in routes)
+
+def test_get_app(internal_routes):
+    app = internal_routes.get_app()
+    assert isinstance(app, web.Application)
+    assert internal_routes._app is not None
+
+def test_get_app_reuse(internal_routes):
+    app1 = internal_routes.get_app()
+    app2 = internal_routes.get_app()
+    assert app1 is app2
+
+@pytest.mark.asyncio
+async def test_routes_added_to_app(aiohttp_client_factory, internal_routes):
+    client = await aiohttp_client_factory()
+    try:
+        resp = await client.get('/files')
+        print(f"Response received: status {resp.status}")
+    except Exception as e:
+        print(f"Exception occurred during GET request: {e}")
+        raise
+
+    assert resp.status != 404, "Route /files does not exist"
+
+@pytest.mark.asyncio
+async def test_file_service_initialization():
+    with patch('api_server.routes.internal.internal_routes.FileService') as MockFileService:
+        # Create a mock instance
+        mock_file_service_instance = MagicMock(spec=FileService)
+        MockFileService.return_value = mock_file_service_instance
+        internal_routes = InternalRoutes()
+
+        # Check if FileService was initialized with the correct parameters
+        MockFileService.assert_called_once_with({
+            "models": models_dir,
+            "user": user_directory,
+            "output": output_directory
+        })
+
+        # Verify that the file_service attribute of InternalRoutes is set
+        assert internal_routes.file_service == mock_file_service_instance

tests-unit/server/services/file_service_test.py

@@ -0,0 +1,54 @@
+import pytest
+from unittest.mock import MagicMock
+from api_server.services.file_service import FileService
+
+@pytest.fixture
+def mock_file_system_ops():
+    return MagicMock()
+
+@pytest.fixture
+def file_service(mock_file_system_ops):
+    allowed_directories = {
+        "models": "/path/to/models",
+        "user": "/path/to/user",
+        "output": "/path/to/output"
+    }
+    return FileService(allowed_directories, file_system_ops=mock_file_system_ops)
+
+def test_list_files_valid_directory(file_service, mock_file_system_ops):
+    mock_file_system_ops.walk_directory.return_value = [
+        {"name": "file1.txt", "path": "file1.txt", "type": "file", "size": 100},
+        {"name": "dir1", "path": "dir1", "type": "directory"}
+    ]
+    
+    result = file_service.list_files("models")
+    
+    assert len(result) == 2
+    assert result[0]["name"] == "file1.txt"
+    assert result[1]["name"] == "dir1"
+    mock_file_system_ops.walk_directory.assert_called_once_with("/path/to/models")
+
+def test_list_files_invalid_directory(file_service):
+    # Does not support walking directories outside of the allowed directories
+    with pytest.raises(ValueError, match="Invalid directory key"):
+        file_service.list_files("invalid_key")
+
+def test_list_files_empty_directory(file_service, mock_file_system_ops):
+    mock_file_system_ops.walk_directory.return_value = []
+    
+    result = file_service.list_files("models")
+    
+    assert len(result) == 0
+    mock_file_system_ops.walk_directory.assert_called_once_with("/path/to/models")
+
+@pytest.mark.parametrize("directory_key", ["models", "user", "output"])
+def test_list_files_all_allowed_directories(file_service, mock_file_system_ops, directory_key):
+    mock_file_system_ops.walk_directory.return_value = [
+        {"name": f"file_{directory_key}.txt", "path": f"file_{directory_key}.txt", "type": "file", "size": 100}
+    ]
+    
+    result = file_service.list_files(directory_key)
+    
+    assert len(result) == 1
+    assert result[0]["name"] == f"file_{directory_key}.txt"
+    mock_file_system_ops.walk_directory.assert_called_once_with(f"/path/to/{directory_key}")

tests-unit/server/utils/file_operations_test.py

@@ -0,0 +1,42 @@
+import pytest
+from typing import List
+from api_server.utils.file_operations import FileSystemOperations, FileSystemItem, is_file_info
+
+@pytest.fixture
+def temp_directory(tmp_path):
+    # Create a temporary directory structure
+    dir1 = tmp_path / "dir1"
+    dir2 = tmp_path / "dir2"
+    dir1.mkdir()
+    dir2.mkdir()
+    (dir1 / "file1.txt").write_text("content1")
+    (dir2 / "file2.txt").write_text("content2")
+    (tmp_path / "file3.txt").write_text("content3")
+    return tmp_path
+
+def test_walk_directory(temp_directory):
+    result: List[FileSystemItem] = FileSystemOperations.walk_directory(str(temp_directory))
+    
+    assert len(result) == 5  # 2 directories and 3 files
+    
+    files = [item for item in result if item['type'] == 'file']
+    dirs = [item for item in result if item['type'] == 'directory']
+    
+    assert len(files) == 3
+    assert len(dirs) == 2
+    
+    file_names = {file['name'] for file in files}
+    assert file_names == {'file1.txt', 'file2.txt', 'file3.txt'}
+    
+    dir_names = {dir['name'] for dir in dirs}
+    assert dir_names == {'dir1', 'dir2'}
+
+def test_walk_directory_empty(tmp_path):
+    result = FileSystemOperations.walk_directory(str(tmp_path))
+    assert len(result) == 0
+
+def test_walk_directory_file_size(temp_directory):
+    result: List[FileSystemItem] = FileSystemOperations.walk_directory(str(temp_directory))
+    files = [item for item in result if is_file_info(item)]
+    for file in files:
+        assert file['size'] > 0  # Assuming all files have some content