diff --git a/README.md b/README.md index 9b4c178..9f45a9c 100644 --- a/README.md +++ b/README.md @@ -158,6 +158,7 @@ These nodes are very lightweight and require no additional dependencies. - is_absolute - Checks if a path is absolute (begins at root directory) - get_size - Returns the size of a file in bytes - get_extension - Extracts the file extension from a path (including the dot) + - set_extension - Replaces or adds a file extension to a path - Directory operations: - list_dir - Lists files and directories in a specified path with filtering options - get_cwd - Returns the current working directory diff --git a/src/basic_data_handling/path_nodes.py b/src/basic_data_handling/path_nodes.py index 3846f1f..a7af9e9 100644 --- a/src/basic_data_handling/path_nodes.py +++ b/src/basic_data_handling/path_nodes.py @@ -453,6 +453,37 @@ class PathNormalize(ComfyNodeABC): return (os.path.normpath(path),) +class PathSetExtension(ComfyNodeABC): + """ + Sets the file extension for a path. + + This node replaces the current extension in a path with a new one. + The extension should include the dot (e.g., '.jpg'). + """ + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "path": (IO.STRING, {}), + "extension": (IO.STRING, {"default": ".txt"}), + } + } + + RETURN_TYPES = (IO.STRING,) + RETURN_NAMES = ("path",) + CATEGORY = "Basic/Path" + DESCRIPTION = cleandoc(__doc__ or "") + FUNCTION = "set_extension" + + def set_extension(self, path: str, extension: str) -> tuple[str]: + # Make sure extension starts with a dot + if not extension.startswith('.') and extension: + extension = '.' + extension + + root, _ = os.path.splitext(path) + return (root + extension,) + + class PathRelative(ComfyNodeABC): """ Returns a relative path. @@ -542,6 +573,7 @@ NODE_CLASS_MAPPINGS = { "Basic data handling: PathExpandVars": PathExpandVars, "Basic data handling: PathGetCwd": PathGetCwd, "Basic data handling: PathGetExtension": PathGetExtension, + "Basic data handling: PathSetExtension": PathSetExtension, "Basic data handling: PathGetSize": PathGetSize, "Basic data handling: PathGlob": PathGlob, "Basic data handling: PathIsAbsolute": PathIsAbsolute, @@ -564,6 +596,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "Basic data handling: PathExpandVars": "expand vars", "Basic data handling: PathGetCwd": "get current working directory", "Basic data handling: PathGetExtension": "get extension", + "Basic data handling: PathSetExtension": "set extension", "Basic data handling: PathGetSize": "get size", "Basic data handling: PathGlob": "glob", "Basic data handling: PathIsAbsolute": "is absolute", diff --git a/tests/test_path_nodes.py b/tests/test_path_nodes.py index 6f68ce9..0052253 100644 --- a/tests/test_path_nodes.py +++ b/tests/test_path_nodes.py @@ -5,7 +5,7 @@ import platform from src.basic_data_handling.path_nodes import ( PathJoin, PathAbspath, PathExists, PathIsFile, PathIsDir, PathGetSize, PathSplit, PathSplitExt, PathBasename, PathDirname, PathGetExtension, - PathNormalize, PathRelative, PathGlob, PathExpandVars, PathGetCwd, + PathSetExtension, PathNormalize, PathRelative, PathGlob, PathExpandVars, PathGetCwd, PathListDir, PathIsAbsolute, PathCommonPrefix, ) @@ -199,6 +199,24 @@ def test_path_get_extension(): # Test with empty string assert node.get_extension("") == ("",) +def test_path_set_extension(): + node = PathSetExtension() + # Test basic extension replacement + assert node.set_extension("file.txt", ".jpg") == ("file.jpg",) + # Test adding extension to file without extension + assert node.set_extension("file", ".png") == ("file.png",) + # Test with extension without dot prefix + assert node.set_extension("document.docx", "pdf") == ("document.pdf",) + # Test with path containing directory + assert node.set_extension("folder/image.jpg", ".png") == ("folder/image.png",) + # Test with multiple dots in filename + assert node.set_extension("archive.tar.gz", ".zip") == ("archive.tar.zip",) + # Test with empty extension (removes extension) + assert node.set_extension("config.ini", "") == ("config",) + # Test with empty string as path + assert node.set_extension("", ".txt") == (".txt",) + + def test_path_normalize(): node = PathNormalize()