use typeddicts for getuioutputs return type

comfy_api/latest/_io.py

@@ -24,6 +24,7 @@ if TYPE_CHECKING:
     from comfy.sd import CLIP, VAE
     from comfy.sd import StyleModel as StyleModel_
     from comfy_api.input import VideoInput
+    from ._ui_types import UIOutputDict
 from comfy_api.internal import (_ComfyNodeInternal, _NodeOutputInternal, classproperty, copy_class, first_real_override, is_class,
     prune_dict, shallow_clone_class)
 from comfy_execution.graph_utils import ExecutionBlocker
@@ -1873,7 +1874,7 @@ class _UIOutput(ABC):
         pass
 
     @abstractmethod
-    def as_dict(self) -> dict:
+    def as_dict(self) -> "UIOutputDict":
         ...
 
 

comfy_api/latest/_ui.py

@@ -5,6 +5,7 @@ import os
 import random
 import uuid
 from io import BytesIO
+from typing import TYPE_CHECKING, Any
 
 import av
 import numpy as np
@@ -23,6 +24,15 @@ import folder_paths
 from comfy.cli_args import args
 from ._io import ComfyNode, FolderType, Image, _UIOutput
 
+if TYPE_CHECKING:
+    from ._ui_types import (
+        ImagesUIOutput,
+        AudioUIOutput,
+        VideoUIOutput,
+        TextUIOutput,
+        UI3DUIOutput,
+    )
+
 
 class SavedResult(dict):
     def __init__(self, filename: str, subfolder: str, type: FolderType):
@@ -48,11 +58,11 @@ class SavedImages(_UIOutput):
         self.results = results
         self.is_animated = is_animated
 
-    def as_dict(self) -> dict:
+    def as_dict(self) -> "ImagesUIOutput":
-        data = {"images": self.results}
+        data: dict[str, Any] = {"images": self.results}
         if self.is_animated:
             data["animated"] = (True,)
-        return data
+        return data  # type: ignore[return-value]
 
 
 class SavedAudios(_UIOutput):
@@ -61,8 +71,8 @@ class SavedAudios(_UIOutput):
         super().__init__()
         self.results = results
 
-    def as_dict(self) -> dict:
+    def as_dict(self) -> "AudioUIOutput":
-        return {"audio": self.results}
+        return {"audio": self.results}  # type: ignore[return-value]
 
 
 def _get_directory_by_folder_type(folder_type: FolderType) -> str:
@@ -397,9 +407,9 @@ class PreviewImage(_UIOutput):
         )
         self.animated = animated
 
-    def as_dict(self):
+    def as_dict(self) -> "ImagesUIOutput":
         return {
-            "images": self.values,
+            "images": self.values,  # type: ignore[typeddict-item]
             "animated": (self.animated,)
         }
 
@@ -421,16 +431,16 @@ class PreviewAudio(_UIOutput):
             quality="128k",
         )
 
-    def as_dict(self) -> dict:
+    def as_dict(self) -> "AudioUIOutput":
-        return {"audio": self.values}
+        return {"audio": self.values}  # type: ignore[return-value]
 
 
 class PreviewVideo(_UIOutput):
     def __init__(self, values: list[SavedResult | dict], **kwargs):
         self.values = values
 
-    def as_dict(self):
+    def as_dict(self) -> "VideoUIOutput":
-        return {"images": self.values, "animated": (True,)}
+        return {"images": self.values, "animated": (True,)}  # type: ignore[return-value]
 
 
 class PreviewUI3D(_UIOutput):
@@ -448,7 +458,7 @@ class PreviewUI3D(_UIOutput):
             img.save(bg_image_path, compress_level=1)
             self.bg_image_path = f"temp/{filename}"
 
-    def as_dict(self):
+    def as_dict(self) -> "UI3DUIOutput":
         return {"result": [self.model_file, self.camera_info, self.bg_image_path]}
 
 
@@ -456,7 +466,7 @@ class PreviewText(_UIOutput):
     def __init__(self, value: str, **kwargs):
         self.value = value
 
-    def as_dict(self):
+    def as_dict(self) -> "TextUIOutput":
         return {"text": (self.value,)}
 
 

comfy_api/latest/_ui_public.py

@@ -1 +1,2 @@
 from ._ui import *  # noqa: F403
+from ._ui_types import *  # noqa: F403

comfy_api/latest/_ui_types.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from typing import Any, Literal, Union
+from typing_extensions import NotRequired, TypedDict
+
+
+FolderTypeLiteral = Literal["input", "output", "temp"]
+
+
+class FileResultDict(TypedDict):
+    filename: str
+    subfolder: str
+    type: FolderTypeLiteral
+
+
+class ImagesUIOutput(TypedDict):
+    images: list[FileResultDict]
+    animated: NotRequired[tuple[bool]]
+
+
+class AudioUIOutput(TypedDict):
+    audio: list[FileResultDict]
+
+
+class VideoUIOutput(TypedDict):
+    images: list[FileResultDict | dict[str, Any]]
+    animated: tuple[Literal[True]]
+
+
+class TextUIOutput(TypedDict):
+    text: tuple[str, ...]
+
+
+class CameraInfoDict(TypedDict, total=False):
+    position: dict[str, float | int]
+    target: dict[str, float | int]
+    zoom: int
+    cameraType: str
+
+
+class UI3DUIOutput(TypedDict):
+    result: list[str | CameraInfoDict | None]
+
+
+class LatentsUIOutput(TypedDict):
+    latents: list[FileResultDict]
+
+
+UIOutputDict = Union[
+    ImagesUIOutput,
+    AudioUIOutput,
+    VideoUIOutput,
+    TextUIOutput,
+    UI3DUIOutput,
+    LatentsUIOutput,
+]