Add type hint for FileLocator (#6968)

* Add type hint for FileLocator

* nit

Author: huchenlei

comfy/comfy_types/__init__.py

@@ -1,6 +1,6 @@
 import torch
 from typing import Callable, Protocol, TypedDict, Optional, List
-from .node_typing import IO, InputTypeDict, ComfyNodeABC, CheckLazyMixin
+from .node_typing import IO, InputTypeDict, ComfyNodeABC, CheckLazyMixin, FileLocator
 
 
 class UnetApplyFunction(Protocol):
@@ -42,4 +42,5 @@ class UnetParams(TypedDict):
     InputTypeDict.__name__,
     ComfyNodeABC.__name__,
     CheckLazyMixin.__name__,
+    FileLocator.__name__,
 ]

comfy/comfy_types/node_typing.py

@@ -295,3 +295,14 @@ def check_lazy_status(self, **kwargs) -> list[str]:
 
         need = [name for name in kwargs if kwargs[name] is None]
         return need
+
+
+class FileLocator(TypedDict):
+    """Provides type hinting for the file location"""
+
+    filename: str
+    """The filename of the file."""
+    subfolder: str
+    """The subfolder of the file."""
+    type: Literal["input", "output", "temp"]
+    """The root folder of the file."""

comfy_extras/nodes_audio.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import torchaudio
 import torch
 import comfy.model_management
@@ -10,6 +12,7 @@
 import hashlib
 import node_helpers
 from comfy.cli_args import args
+from comfy.comfy_types import FileLocator
 
 class EmptyLatentAudio:
     def __init__(self):
@@ -164,7 +167,7 @@ def INPUT_TYPES(s):
     def save_audio(self, audio, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
         filename_prefix += self.prefix_append
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
-        results = list()
+        results: list[FileLocator] = []
 
         metadata = {}
         if not args.disable_metadata:

comfy_extras/nodes_images.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import nodes
 import folder_paths
 from comfy.cli_args import args
@@ -9,6 +11,8 @@
 import json
 import os
 
+from comfy.comfy_types import FileLocator
+
 MAX_RESOLUTION = nodes.MAX_RESOLUTION
 
 class ImageCrop:
@@ -99,7 +103,7 @@ def save_images(self, images, fps, filename_prefix, lossless, quality, method, n
         method = self.methods.get(method)
         filename_prefix += self.prefix_append
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
-        results = list()
+        results: list[FileLocator] = []
         pil_images = []
         for image in images:
             i = 255. * image.cpu().numpy()

comfy_extras/nodes_video.py

@@ -1,9 +1,12 @@
+from __future__ import annotations
+
 import os
 import av
 import torch
 import folder_paths
 import json
 from fractions import Fraction
+from comfy.comfy_types import FileLocator
 
 
 class SaveWEBM:
@@ -62,7 +65,7 @@ def save_images(self, images, codec, fps, filename_prefix, crf, prompt=None, ext
         container.mux(stream.encode())
         container.close()
 
-        results = [{
+        results: list[FileLocator] = [{
             "filename": file,
             "subfolder": subfolder,
             "type": self.type

nodes.py

@@ -25,7 +25,7 @@
 import comfy.sd
 import comfy.utils
 import comfy.controlnet
-from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
+from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict, FileLocator
 
 import comfy.clip_vision
 
@@ -479,7 +479,7 @@ def save(self, samples, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=No
 
         file = f"{filename}_{counter:05}_.latent"
 
-        results = list()
+        results: list[FileLocator] = []
         results.append({
             "filename": file,
             "subfolder": subfolder,