add type hints to server.py

server.py

@@ -1,7 +1,13 @@
+from __future__ import annotations
+
 import os
 import sys
 import asyncio
 import traceback
+from asyncio import AbstractEventLoop
+from collections.abc import Callable
+from enum import Enum
+from typing import TYPE_CHECKING, Any
 
 import nodes
 import folder_paths
@@ -12,9 +18,14 @@ import json
 import glob
 import struct
 import ssl
-from PIL import Image, ImageOps
+import hashlib
+import PIL
+from PIL import ImageOps
+from PIL.Image import Image
 from PIL.PngImagePlugin import PngInfo
 from io import BytesIO
+from aiohttp.typedefs import Handler, Middleware
+from multidict import MultiDictProxy
 
 import aiohttp
 from aiohttp import web
@@ -30,26 +41,41 @@ from app.user_manager import UserManager
 from model_filemanager import download_model, DownloadModelStatus
 from typing import Optional
 
-class BinaryEventTypes:
+if TYPE_CHECKING:
+    from execution import PromptQueue
+
+class BinaryEventTypes(Enum):
     PREVIEW_IMAGE = 1
     UNENCODED_PREVIEW_IMAGE = 2
 
-async def send_socket_catch_exception(function, message):
+class DirType(Enum):
+    input = 0
+    temp = 1
+    output = 2
+
+    @staticmethod
+    def from_str(name: str | None) -> DirType | None:
+        if name is None:
+            return None
+
+        return DirType[name]
+
+async def send_socket_catch_exception(function: Callable[[Any], Any], message: Any) -> None:
     try:
         await function(message)
     except (aiohttp.ClientError, aiohttp.ClientPayloadError, ConnectionResetError) as err:
         logging.warning("send error: {}".format(err))
 
 @web.middleware
-async def cache_control(request: web.Request, handler):
+async def cache_control(request: web.Request, handler: Handler) -> web.StreamResponse:
-    response: web.Response = await handler(request)
+    response = await handler(request)
     if request.path.endswith('.js') or request.path.endswith('.css'):
         response.headers.setdefault('Cache-Control', 'no-cache')
     return response
 
-def create_cors_middleware(allowed_origin: str):
+def create_cors_middleware(allowed_origin: str) -> Middleware:
     @web.middleware
-    async def cors_middleware(request: web.Request, handler):
+    async def cors_middleware(request: web.Request, handler: Handler) -> web.Response:
         if request.method == "OPTIONS":
             # Pre-flight request. Reply successfully:
             response = web.Response()
@@ -65,7 +91,7 @@ def create_cors_middleware(allowed_origin: str):
     return cors_middleware
 
 class PromptServer():
-    def __init__(self, loop):
+    def __init__(self, loop: AbstractEventLoop):
         PromptServer.instance = self
 
         mimetypes.init()
@@ -73,20 +99,20 @@ class PromptServer():
 
         self.user_manager = UserManager()
         self.supports = ["custom_nodes_from_web"]
-        self.prompt_queue = None
+        self.prompt_queue: PromptQueue | None = None
         self.loop = loop
         self.messages = asyncio.Queue()
         self.client_session:Optional[aiohttp.ClientSession] = None
         self.number = 0
 
-        middlewares = [cache_control]
+        middlewares: list[Middleware] = [cache_control]
         if args.enable_cors_header:
             middlewares.append(create_cors_middleware(args.enable_cors_header))
 
         max_upload_size = round(args.max_upload_size * 1024 * 1024)
         self.app = web.Application(client_max_size=max_upload_size, middlewares=middlewares)
-        self.sockets = dict()
+        self.sockets: dict[str, web.WebSocketResponse] = dict()
-        self.web_root = (
+        self.web_root: str = (
             FrontendManager.init_frontend(args.front_end_version)
             if args.front_end_root is None
             else args.front_end_root
@@ -95,12 +121,13 @@ class PromptServer():
         routes = web.RouteTableDef()
         self.routes = routes
         self.last_node_id = None
-        self.client_id = None
+        self.client_id: str | None = None
+        self.last_prompt_id: str | None = None
 
         self.on_prompt_handlers = []
 
         @routes.get('/ws')
-        async def websocket_handler(request):
+        async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
             ws = web.WebSocketResponse()
             await ws.prepare(request)
             sid = request.rel_url.query.get('clientId', '')
@@ -127,7 +154,7 @@ class PromptServer():
             return ws
 
         @routes.get("/")
-        async def get_root(request):
+        async def get_root(request: web.Request) -> web.FileResponse:
             response = web.FileResponse(os.path.join(self.web_root, "index.html"))
             response.headers['Cache-Control'] = 'no-cache'
             response.headers["Pragma"] = "no-cache"
@@ -135,12 +162,12 @@ class PromptServer():
             return response
 
         @routes.get("/embeddings")
-        def get_embeddings(self):
+        def get_embeddings(request: web.Request) -> web.Response:
             embeddings = folder_paths.get_filename_list("embeddings")
             return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings)))
 
         @routes.get("/extensions")
-        async def get_extensions(request):
+        async def get_extensions(request: web.Request) -> web.Response:
             files = glob.glob(os.path.join(
                 glob.escape(self.web_root), 'extensions/**/*.js'), recursive=True)
 
@@ -153,20 +180,22 @@ class PromptServer():
 
             return web.json_response(extensions)
 
-        def get_dir_by_type(dir_type):
+        def get_dir_by_type(dir_type: DirType | None) -> tuple[str, DirType]:
             if dir_type is None:
-                dir_type = "input"
+                dir_type = DirType.input
 
-            if dir_type == "input":
+            if dir_type == DirType.input:
                 type_dir = folder_paths.get_input_directory()
-            elif dir_type == "temp":
+            elif dir_type == DirType.temp:
                 type_dir = folder_paths.get_temp_directory()
-            elif dir_type == "output":
+            elif dir_type == DirType.output:
                 type_dir = folder_paths.get_output_directory()
+            else:
+                raise ValueError("Unknown dir type")
 
             return type_dir, dir_type
 
-        def compare_image_hash(filepath, image):
+        def compare_image_hash(filepath: str, image: web.FileField) -> bool:
             hasher = node_helpers.hasher()
             
             # function to compare hashes of two images to see if it already exists, fix to #3465
@@ -181,12 +210,14 @@ class PromptServer():
                 return a.hexdigest() == b.hexdigest()
             return False
 
-        def image_upload(post, image_save_function=None):
+        def image_upload(
+                post: MultiDictProxy[str | bytes | web.FileField],
+                image_save_function: Callable[[web.FileField, MultiDictProxy[str | bytes | web.FileField], str], web.Response] | None=None) -> web.Response:
             image = post.get("image")
             overwrite = post.get("overwrite")
             image_is_duplicate = False
 
-            image_upload_type = post.get("type")
+            image_upload_type = DirType.from_str(post.get("type"))
             upload_dir, image_upload_type = get_dir_by_type(image_upload_type)
 
             if image and image.file:
@@ -225,21 +256,21 @@ class PromptServer():
                         with open(filepath, "wb") as f:
                             f.write(image.file.read())
 
-                return web.json_response({"name" : filename, "subfolder": subfolder, "type": image_upload_type})
+                return web.json_response({"name" : filename, "subfolder": subfolder, "type": image_upload_type.name})
             else:
                 return web.Response(status=400)
 
         @routes.post("/upload/image")
-        async def upload_image(request):
+        async def upload_image(request: web.Request) -> web.Response:
             post = await request.post()
             return image_upload(post)
 
 
         @routes.post("/upload/mask")
-        async def upload_mask(request):
+        async def upload_mask(request: web.Request) -> web.Response:
             post = await request.post()
 
-            def image_save_function(image, post, filepath):
+            def image_save_function(image: web.FileField, post: MultiDictProxy[str | bytes | web.FileField], filepath: str) -> web.Response:
                 original_ref = json.loads(post.get("original_ref"))
                 filename, output_dir = folder_paths.annotated_filepath(original_ref['filename'])
 
@@ -263,13 +294,13 @@ class PromptServer():
                 file = os.path.join(output_dir, filename)
 
                 if os.path.isfile(file):
-                    with Image.open(file) as original_pil:
+                    with PIL.Image.open(file) as original_pil:
                         metadata = PngInfo()
                         if hasattr(original_pil,'text'):
                             for key in original_pil.text:
                                 metadata.add_text(key, original_pil.text[key])
                         original_pil = original_pil.convert('RGBA')
-                        mask_pil = Image.open(image.file).convert('RGBA')
+                        mask_pil = PIL.Image.open(image.file).convert('RGBA')
 
                         # alpha copy
                         new_alpha = mask_pil.getchannel('A')
@@ -279,7 +310,7 @@ class PromptServer():
             return image_upload(post, image_save_function)
 
         @routes.get("/view")
-        async def view_image(request):
+        async def view_image(request: web.Request) -> web.StreamResponse:
             if "filename" in request.rel_url.query:
                 filename = request.rel_url.query["filename"]
                 filename,output_dir = folder_paths.annotated_filepath(filename)
@@ -306,7 +337,7 @@ class PromptServer():
 
                 if os.path.isfile(file):
                     if 'preview' in request.rel_url.query:
-                        with Image.open(file) as img:
+                        with PIL.Image.open(file) as img:
                             preview_info = request.rel_url.query['preview'].split(';')
                             image_format = preview_info[0]
                             if image_format not in ['webp', 'jpeg'] or 'a' in request.rel_url.query.get('channel', ''):
@@ -331,10 +362,10 @@ class PromptServer():
                         channel = request.rel_url.query["channel"]
 
                     if channel == 'rgb':
-                        with Image.open(file) as img:
+                        with PIL.Image.open(file) as img:
                             if img.mode == "RGBA":
                                 r, g, b, a = img.split()
-                                new_img = Image.merge('RGB', (r, g, b))
+                                new_img = PIL.Image.merge('RGB', (r, g, b))
                             else:
                                 new_img = img.convert("RGB")
 
@@ -346,14 +377,14 @@ class PromptServer():
                                                 headers={"Content-Disposition": f"filename=\"{filename}\""})
 
                     elif channel == 'a':
-                        with Image.open(file) as img:
+                        with PIL.Image.open(file) as img:
                             if img.mode == "RGBA":
                                 _, _, _, a = img.split()
                             else:
-                                a = Image.new('L', img.size, 255)
+                                a = PIL.Image.new('L', img.size, 255)
 
                             # alpha img
-                            alpha_img = Image.new('RGBA', img.size)
+                            alpha_img = PIL.Image.new('RGBA', img.size)
                             alpha_img.putalpha(a)
                             alpha_buffer = BytesIO()
                             alpha_img.save(alpha_buffer, format='PNG')
@@ -367,8 +398,8 @@ class PromptServer():
             return web.Response(status=404)
 
         @routes.get("/view_metadata/{folder_name}")
-        async def view_metadata(request):
+        async def view_metadata(request: web.Request) -> web.Response:
-            folder_name = request.match_info.get("folder_name", None)
+            folder_name: str | None = request.match_info.get("folder_name", None)
             if folder_name is None:
                 return web.Response(status=404)
             if not "filename" in request.rel_url.query:
@@ -390,7 +421,7 @@ class PromptServer():
             return web.json_response(dt["__metadata__"])
 
         @routes.get("/system_stats")
-        async def get_queue(request):
+        async def get_queue(request: web.Request) -> web.Response:
             device = comfy.model_management.get_torch_device()
             device_name = comfy.model_management.get_torch_device_name(device)
             vram_total, torch_vram_total = comfy.model_management.get_total_memory(device, torch_total_too=True)
@@ -416,7 +447,7 @@ class PromptServer():
             return web.json_response(system_stats)
 
         @routes.get("/prompt")
-        async def get_prompt(request):
+        async def get_prompt(request: web.Request) -> web.Response:
             return web.json_response(self.get_queue_info())
 
         def node_info(node_class):
@@ -444,7 +475,7 @@ class PromptServer():
             return info
 
         @routes.get("/object_info")
-        async def get_object_info(request):
+        async def get_object_info(request: web.Request) -> web.Response:
             out = {}
             for x in nodes.NODE_CLASS_MAPPINGS:
                 try:
@@ -455,7 +486,7 @@ class PromptServer():
             return web.json_response(out)
 
         @routes.get("/object_info/{node_class}")
-        async def get_object_info_node(request):
+        async def get_object_info_node(request: web.Request) -> web.Response:
             node_class = request.match_info.get("node_class", None)
             out = {}
             if (node_class is not None) and (node_class in nodes.NODE_CLASS_MAPPINGS):
@@ -463,19 +494,19 @@ class PromptServer():
             return web.json_response(out)
 
         @routes.get("/history")
-        async def get_history(request):
+        async def get_history(request: web.Request) -> web.Response:
             max_items = request.rel_url.query.get("max_items", None)
             if max_items is not None:
                 max_items = int(max_items)
             return web.json_response(self.prompt_queue.get_history(max_items=max_items))
 
         @routes.get("/history/{prompt_id}")
-        async def get_history(request):
+        async def get_history(request: web.Request) -> web.Response:
             prompt_id = request.match_info.get("prompt_id", None)
             return web.json_response(self.prompt_queue.get_history(prompt_id=prompt_id))
 
         @routes.get("/queue")
-        async def get_queue(request):
+        async def get_queue(request: web.Request) -> web.Response:
             queue_info = {}
             current_queue = self.prompt_queue.get_current_queue()
             queue_info['queue_running'] = current_queue[0]
@@ -483,7 +514,7 @@ class PromptServer():
             return web.json_response(queue_info)
 
         @routes.post("/prompt")
-        async def post_prompt(request):
+        async def post_prompt(request: web.Request) -> web.Response:
             logging.info("got prompt")
             resp_code = 200
             out_string = ""
@@ -522,7 +553,7 @@ class PromptServer():
                 return web.json_response({"error": "no prompt", "node_errors": []}, status=400)
 
         @routes.post("/queue")
-        async def post_queue(request):
+        async def post_queue(request: web.Request) -> web.Response:
             json_data =  await request.json()
             if "clear" in json_data:
                 if json_data["clear"]:
@@ -536,12 +567,12 @@ class PromptServer():
             return web.Response(status=200)
 
         @routes.post("/interrupt")
-        async def post_interrupt(request):
+        async def post_interrupt(request: web.Request) -> web.Response:
             nodes.interrupt_processing()
             return web.Response(status=200)
 
         @routes.post("/free")
-        async def post_free(request):
+        async def post_free(request: web.Request) -> web.Response:
             json_data = await request.json()
             unload_models = json_data.get("unload_models", False)
             free_memory = json_data.get("free_memory", False)
@@ -552,7 +583,7 @@ class PromptServer():
             return web.Response(status=200)
 
         @routes.post("/history")
-        async def post_history(request):
+        async def post_history(request: web.Request) -> web.Response:
             json_data =  await request.json()
             if "clear" in json_data:
                 if json_data["clear"]:
@@ -584,7 +615,7 @@ class PromptServer():
             if session is None:
                 logging.error("Client session is not initialized")
                 return web.Response(status=500)
-            
+
             task = asyncio.create_task(download_model(lambda url: session.get(url), model_filename, url, model_directory, report_progress, progress_interval))
             await task
 
@@ -594,7 +625,7 @@ class PromptServer():
         timeout = aiohttp.ClientTimeout(total=None) # no timeout
         self.client_session = aiohttp.ClientSession(timeout=timeout)
 
-    def add_routes(self):
+    def add_routes(self) -> None:
         self.user_manager.add_routes(self.routes)
 
         # Prefix every route with /api for easier matching for delegation.
@@ -620,14 +651,14 @@ class PromptServer():
             web.static('/', self.web_root),
         ])
 
-    def get_queue_info(self):
+    def get_queue_info(self) -> dict[str, any]:
         prompt_info = {}
         exec_info = {}
         exec_info['queue_remaining'] = self.prompt_queue.get_tasks_remaining()
         prompt_info['exec_info'] = exec_info
         return prompt_info
 
-    async def send(self, event, data, sid=None):
+    async def send(self, event: BinaryEventTypes | str, data: tuple[str, Image, int] | bytes | bytearray | dict, sid: str | None=None):
         if event == BinaryEventTypes.UNENCODED_PREVIEW_IMAGE:
             await self.send_image(data, sid=sid)
         elif isinstance(data, (bytes, bytearray)):
@@ -635,24 +666,24 @@ class PromptServer():
         else:
             await self.send_json(event, data, sid)
 
-    def encode_bytes(self, event, data):
+    def encode_bytes(self, event: BinaryEventTypes, data: bytes | bytearray):
-        if not isinstance(event, int):
+        if not isinstance(event, BinaryEventTypes):
-            raise RuntimeError(f"Binary event types must be integers, got {event}")
+            raise RuntimeError(f"Binary event types must be BinaryEventTypes, got {event}")
 
         packed = struct.pack(">I", event)
         message = bytearray(packed)
         message.extend(data)
         return message
 
-    async def send_image(self, image_data, sid=None):
+    async def send_image(self, image_data: tuple[str, Image, int] | None, sid: str | None=None) -> None:
-        image_type = image_data[0]
+        image_type: str = image_data[0]
-        image = image_data[1]
+        image: Image = image_data[1]
-        max_size = image_data[2]
+        max_size: int = image_data[2]
         if max_size is not None:
-            if hasattr(Image, 'Resampling'):
+            if hasattr(PIL.Image, 'Resampling'):
-                resampling = Image.Resampling.BILINEAR
+                resampling = PIL.Image.Resampling.BILINEAR
             else:
-                resampling = Image.ANTIALIAS
+                resampling = PIL.Image.ANTIALIAS
 
             image = ImageOps.contain(image, (max_size, max_size), resampling)
         type_num = 1
@@ -668,7 +699,7 @@ class PromptServer():
         preview_bytes = bytesIO.getvalue()
         await self.send_bytes(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes, sid=sid)
 
-    async def send_bytes(self, event, data, sid=None):
+    async def send_bytes(self, event: BinaryEventTypes, data: bytes | bytearray, sid: str | None=None) -> None:
         message = self.encode_bytes(event, data)
 
         if sid is None:
@@ -678,7 +709,7 @@ class PromptServer():
         elif sid in self.sockets:
             await send_socket_catch_exception(self.sockets[sid].send_bytes, message)
 
-    async def send_json(self, event, data, sid=None):
+    async def send_json(self, event, data: dict, sid: str | None=None) -> None:
         message = {"type": event, "data": data}
 
         if sid is None:
@@ -688,19 +719,19 @@ class PromptServer():
         elif sid in self.sockets:
             await send_socket_catch_exception(self.sockets[sid].send_json, message)
 
-    def send_sync(self, event, data, sid=None):
+    def send_sync(self, event: BinaryEventTypes | str, data: tuple[str, Image, int] | bytes | bytearray | dict, sid: str | None=None) -> None:
         self.loop.call_soon_threadsafe(
             self.messages.put_nowait, (event, data, sid))
 
-    def queue_updated(self):
+    def queue_updated(self) -> None:
         self.send_sync("status", { "status": self.get_queue_info() })
 
-    async def publish_loop(self):
+    async def publish_loop(self) -> None:
         while True:
             msg = await self.messages.get()
             await self.send(*msg)
 
-    async def start(self, address, port, verbose=True, call_on_start=None):
+    async def start(self, address: str, port: int, verbose=True, call_on_start: Callable[[str, str, int], None] | None=None):
         runner = web.AppRunner(self.app, access_log=None)
         await runner.setup()
         ssl_ctx = None