ChrisColeTech-ComfyUI-Line-counter/line_counter.py

import os
import hashlib
import random


class DirectoryFileCounter():
    SEARCH_ALIASES = ['count files', 'file count', 'directory count', 'number of files']
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "directory_path": ("STRING", {"default": ""}),
            },
        }

    RETURN_TYPES = ("INT", "FLOAT", "STRING",)

    FUNCTION = "execute_file_counter"

    CATEGORY = "🤖 CCTech/Files"

    def execute_file_counter(self, directory_path: str) -> dict:
        # Expand environment variables in the directory path
        directory_path = os.path.expandvars(directory_path)

        if not os.path.isdir(directory_path):
            raise FileNotFoundError(
                f"The path '{directory_path}' does not exist.")

        file_count = 0

        # Walk the directory tree
        for root, dirs, files in os.walk(directory_path):
            file_count += len(files)  # Count files in each directory

        count = file_count - 1
        return (count, float(count),  str(count))


class TextFileLineCounter():
    SEARCH_ALIASES = ['count lines', 'line count', 'text file lines']
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "file_path": ("STRING", {"default": ""}),
            },
        }

    RETURN_TYPES = ("INT", "FLOAT", "STRING",)

    FUNCTION = "execute_line_counter"

    CATEGORY = "🤖 CCTech/Files"

    def execute_line_counter(self, file_path: str) -> dict:
        # Expand environment variables in the file path
        file_path = os.path.expandvars(file_path)

        if not os.path.isfile(file_path):
            raise FileNotFoundError(
                f"The file at path '{file_path}' does not exist.")
        with open(file_path, 'r', encoding='utf-8') as file:
            lines = file.readlines()
        line_count = len([line for line in lines if line.strip() != ""]) - 1
        return (line_count, float(line_count),  str(line_count))


class SimpleNumberCounter:
    SEARCH_ALIASES = ['counter', 'increment', 'decrement', 'randomize', 'number counter', 'index counter']
    def __init__(self):
        self.counters = {}

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "reset_bool": ("BOOLEAN", {"default": False}),
                "number_type": (["integer", "float"],),
                "mode": (["increment", "decrement", "increment_to_stop", "decrement_to_stop", "randomize"],),
                "start": ("FLOAT", {"default": 0, "min": -18446744073709551615, "max": 18446744073709551615, "step": 0.01}),
                "stop": ("FLOAT", {"default": 0, "min": -18446744073709551615, "max": 18446744073709551615, "step": 0.01}),
                "step": ("FLOAT", {"default": 1, "min": 0, "max": 99999, "step": 0.01}),
            },
            "hidden": {
                "unique_id": "UNIQUE_ID",
            }
        }

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

    RETURN_TYPES = ("NUMBER", "FLOAT", "INT")
    RETURN_NAMES = ("number", "float", "int")
    FUNCTION = "increment_number"

    CATEGORY = "🤖 CCTech/Utilities"

    def increment_number(self, reset_bool, number_type, mode, start, stop, step, unique_id):

        counter = int(start) if number_type == 'integer' else start
        if unique_id in self.counters:
            counter = self.counters[unique_id]

        if reset_bool:
            counter = start
        elif mode == 'increment':
            counter += step
        elif mode == 'decrement':
            counter -= step
        elif mode == 'increment_to_stop':
            counter = counter + step if counter < stop else counter
        elif mode == 'decrement_to_stop':
            counter = counter - step if counter > stop else counter
        elif mode == 'randomize':
            lo, hi = (start, stop) if start <= stop else (stop, start)
            if number_type == 'integer':
                a, b = int(lo), int(hi)
                if a > b:
                    a, b = b, a
                counter = random.randint(a, b)
            else:
                counter = random.uniform(lo, hi)

        self.counters[unique_id] = counter

        result = int(counter) if number_type == 'integer' else float(counter)

        return (result, float(counter), int(counter))


TEXT_TYPE = "STRING"


class TextFileLineReader:
    SEARCH_ALIASES = ['read line', 'read text file', 'sequential text', 'line reader', 'read prompts']
    # Class-level variable to keep track of the current line index
    current_index = 0

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "file_path": ("STRING", {"default": ''}),
                # Add mode input
                "mode": (["automatic", "index"], {"default": "automatic"}),
                # Add index input
                "index": ("INT", {"default": 0, "min": 0, "step": 1}),
            },
        }

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

    RETURN_TYPES = (TEXT_TYPE,)
    RETURN_NAMES = ("line_text",)
    FUNCTION = "read_line"

    CATEGORY = "🤖 CCTech/Files"

    def read_line(self, file_path: str, mode: str, index: int):
        # Expand environment variables in the file path
        file_path = os.path.expandvars(file_path)

        # Check if the file exists
        if not os.path.isfile(file_path):
            raise FileNotFoundError(
                f"The file at path '{file_path}' does not exist.")

        # Read all lines from the file
        with open(file_path, 'r', encoding='utf-8') as file:
            lines = [line.strip()
                     for line in file if line.strip()]  # Skip empty lines

        # If there are no lines, return an empty string
        if not lines:
            return ("",)

        # Handle the mode and index
        if mode == "index":
            # Use the provided index directly
            if index < 0 or index >= len(lines):
                print(f"Index `{index}` is out of bounds. Resetting to 0.")
                index = 0  # Reset to 0 if out of bounds
            line_text = lines[index]
        else:
            # Automatic mode: use the current_index and increment it
            if self.current_index >= len(lines):
                print("Reached the end of the file. Resetting to 0.")
                self.current_index = 0  # Reset to 0 if at the end
            line_text = lines[self.current_index]
            self.current_index += 1  # Increment the index for the next run

        return (line_text,)


NODE_CLASS_MAPPINGS = {
    "Text File Line Counter": TextFileLineCounter,
    "Directory File Counter": DirectoryFileCounter,
    "Simple Number Counter": SimpleNumberCounter,
    "Text File Line Reader": TextFileLineReader
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Text File Line Counter": "Text File Line Counter ⏲️",
    "Directory File Counter": "Directory File Counter ⏲️",
    "Simple Number Counter": "Simple Number Counter ⏲️",
    "Text File Line Reader": "Text File Line Reader ⏲️"
}