From 5d84607bf3a761d796fb0cf3b6fdba8480ead5f7 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Wed, 5 Mar 2025 15:35:26 -0500
Subject: [PATCH] Add type hint for FileLocator (#6968)

* Add type hint for FileLocator

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

diff --git a/comfy/comfy_types/__init__.py b/comfy/comfy_types/__init__.py
index 19ec33f9844..7640fbe3f82 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 @@ class UnetParams(TypedDict):
     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 6146b70f86b..fe130567da5 100644
--- a/comfy/comfy_types/node_typing.py
+++ b/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."""
diff --git a/comfy_extras/nodes_audio.py b/comfy_extras/nodes_audio.py
index 3cb918e09d3..136ad6159b8 100644
--- a/comfy_extras/nodes_audio.py
+++ b/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:
diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py
index af37666b29f..e11a4583a5c 100644
--- a/comfy_extras/nodes_images.py
+++ b/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()
diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py
index 53920ba1828..97ca513d8f4 100644
--- a/comfy_extras/nodes_video.py
+++ b/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
diff --git a/nodes.py b/nodes.py
index dec6cdc86ac..bbf49915cf7 100644
--- a/nodes.py
+++ b/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,