From 446d7be8cc77e0c0c8825b86f0f3974ded18f659 Mon Sep 17 00:00:00 2001
From: huchenlei <huchenlei@proton.me>
Date: Tue, 25 Feb 2025 11:13:38 -0500
Subject: [PATCH] Add type hint for FileLocator

---
 comfy/comfy_types/__init__.py    |  3 ++-
 comfy/comfy_types/node_typing.py | 11 +++++++++++
 comfy_extras/nodes_audio.py      |  3 ++-
 comfy_extras/nodes_images.py     |  4 +++-
 comfy_extras/nodes_video.py      |  3 ++-
 nodes.py                         |  4 ++--
 6 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/comfy/comfy_types/__init__.py b/comfy/comfy_types/__init__.py
index 19ec33f9..7640fbe3 100644
--- a/comfy/comfy_types/__init__.py
+++ b/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 @@ __all__ = [
     InputTypeDict.__name__,
     ComfyNodeABC.__name__,
     CheckLazyMixin.__name__,
+    FileLocator.__name__,
 ]
diff --git a/comfy/comfy_types/node_typing.py b/comfy/comfy_types/node_typing.py
index 0696dbe5..acf91281 100644
--- a/comfy/comfy_types/node_typing.py
+++ b/comfy/comfy_types/node_typing.py
@@ -293,3 +293,14 @@ class CheckLazyMixin:
 
         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."""
diff --git a/comfy_extras/nodes_audio.py b/comfy_extras/nodes_audio.py
index 3cb918e0..7609b35c 100644
--- a/comfy_extras/nodes_audio.py
+++ b/comfy_extras/nodes_audio.py
@@ -10,6 +10,7 @@ import random
 import hashlib
 import node_helpers
 from comfy.cli_args import args
+from comfy.comfy_types import FileLocator
 
 class EmptyLatentAudio:
     def __init__(self):
@@ -164,7 +165,7 @@ class SaveAudio:
     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:
diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py
index af37666b..3503fb99 100644
--- a/comfy_extras/nodes_images.py
+++ b/comfy_extras/nodes_images.py
@@ -9,6 +9,8 @@ import numpy as np
 import json
 import os
 
+from comfy.comfy_types import FileLocator
+
 MAX_RESOLUTION = nodes.MAX_RESOLUTION
 
 class ImageCrop:
@@ -99,7 +101,7 @@ class SaveAnimatedWEBP:
         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()
diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py
index f3922e03..8e7f6b97 100644
--- a/comfy_extras/nodes_video.py
+++ b/comfy_extras/nodes_video.py
@@ -4,6 +4,7 @@ import torch
 import folder_paths
 import json
 from fractions import Fraction
+from comfy.comfy_types import FileLocator
 
 
 class SaveWEBM:
@@ -61,7 +62,7 @@ class SaveWEBM:
                 container.mux(packet)
         container.close()
 
-        results = [{
+        results: list[FileLocator] = [{
             "filename": file,
             "subfolder": subfolder,
             "type": self.type
diff --git a/nodes.py b/nodes.py
index 4e68af79..03036069 100644
--- a/nodes.py
+++ b/nodes.py
@@ -25,7 +25,7 @@ import comfy.sample
 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 @@ class SaveLatent:
 
         file = f"{filename}_{counter:05}_.latent"
 
-        results = list()
+        results: list[FileLocator] = []
         results.append({
             "filename": file,
             "subfolder": subfolder,