List files.

This commit is contained in:
Robin Huang 2024-08-09 12:56:37 -07:00
parent d73dd32be1
commit e4f5e7ac7c

View File

@ -1,14 +1,61 @@
from aiohttp import web from aiohttp import web
from typing import Optional, Dict, List
from folder_paths import models_dir, user_directory, output_directory
import os
class InternalRoutes: class InternalRoutes:
def __init__(self): def __init__(self):
self.routes = web.RouteTableDef() self.routes: web.RouteTableDef = web.RouteTableDef()
self._app = None self._app: Optional[web.Application] = None
self.allowed_directories: Dict[str, str] = {
"models": models_dir,
"user": user_directory,
"output": output_directory
}
def setup_routes(self): def setup_routes(self):
@self.routes.get('/files/') @self.routes.get('/files')
async def internal_test(request): async def list_files(request):
return web.json_response({"message": "Internal route is working"}) directory_key = request.query.get('directory', '')
if directory_key not in self.allowed_directories:
return web.json_response({"error": "Invalid directory key"}, status=400)
directory_path = self.allowed_directories[directory_key]
try:
file_list = self._get_file_list(directory_path)
return web.json_response({"files": file_list})
except FileNotFoundError:
return web.json_response({"error": "Directory not found"}, status=404)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
def _get_file_list(self, directory: str) -> List[Dict[str, Optional[str]]]:
if not os.path.exists(directory):
raise FileNotFoundError(f"Directory not found: {directory}")
file_list = []
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",
"size": None
})
return file_list
def get_app(self): def get_app(self):
if self._app is None: if self._app is None: