kohya_ss/kohya_gui/class_gui_config.py

import toml
from .common_gui import scriptdir
from .custom_logging import setup_logging

# Set up logging
log = setup_logging()


class KohyaSSGUIConfig:
    """
    A class to handle the configuration for the Kohya SS GUI.
    """

    def __init__(self, config_file_path: str = "./config.toml"):
        """
        Initialize the KohyaSSGUIConfig class.
        """
        self.config = self.load_config(config_file_path=config_file_path)
        # Initialize last_used_folder during class initialization
        self.get_last_used_folder()

    def load_config(self, config_file_path: str = "./config.toml") -> dict:
        """
        Loads the Kohya SS GUI configuration from a TOML file.

        Returns:
        dict: The configuration data loaded from the TOML file.
        """
        try:
            # Attempt to load the TOML configuration file from the specified directory.
            config = toml.load(f"{config_file_path}")
            log.debug(f"Loaded configuration from {config_file_path}")
        except FileNotFoundError:
            # If the config file is not found, initialize `config` as an empty dictionary to handle missing configurations gracefully.
            config = {}
            log.debug(
                f"No configuration file found at {config_file_path}. Initializing empty configuration."
            )

        # Ensure last_used_folder is present, defaulting to scriptdir
        if "last_used_folder" not in config:
            config["last_used_folder"] = scriptdir
            log.debug(f"Initialized 'last_used_folder' to '{scriptdir}'")

        return config

    def save_config(self, config: dict, config_file_path: str = "./config.toml"):
        """
        Saves the Kohya SS GUI configuration to a TOML file.

        Parameters:
        - config (dict): The configuration data to save.
        """
        # Write the configuration data to the TOML file
        with open(f"{config_file_path}", "w", encoding="utf-8") as f:
            toml.dump(config, f)

    def get(self, key: str, default=None):
        """
        Retrieves the value of a specified key from the configuration data.

        Parameters:
        - key (str): The key to retrieve the value for.
        - default: The default value to return if the key is not found.

        Returns:
        The value associated with the key, or the default value if the key is not found.
        """
        # Split the key into a list of keys if it contains a dot (.)
        keys = key.split(".")
        # Initialize `data` with the entire configuration data
        data = self.config

        # Iterate over the keys to access nested values
        for k in keys:
            log.debug(k)
            # If the key is not found in the current data, return the default value
            if k not in data:
                log.debug(
                    f"Key '{key}' not found in configuration. Returning default value."
                )
                return default

            # Update `data` to the value associated with the current key
            data = data.get(k)

        # Return the final value
        log.debug(f"Returned {data}")
        return data

    def is_config_loaded(self) -> bool:
        """
        Checks if the configuration was loaded from a file.

        Returns:
        bool: True if the configuration was loaded from a file, False otherwise.
        """
        is_loaded = self.config != {}
        log.debug(f"Configuration was loaded from file: {is_loaded}")
        return is_loaded

    def get_last_used_folder(self) -> str:
        """
        Retrieves the last used folder from the configuration.

        Returns:
        str: The last used folder path.
        """
        folder = self.config.get("last_used_folder", scriptdir)
        # Ensure that the returned path is a string, even if it's empty or None from config
        if not isinstance(folder, str):
            folder = str(folder) if folder is not None else scriptdir
            # Update the config if the type was wrong.
            self.config["last_used_folder"] = folder
        log.debug(f"Retrieved last_used_folder: {folder}")
        return folder

    def set_last_used_folder(self, folder_path: str):
        """
        Sets the last used folder in the configuration.

        Parameters:
        - folder_path (str): The path to the folder to be set.
        """
        if not isinstance(folder_path, str):
            log.error(f"Attempted to set last_used_folder with non-string value: {folder_path}")
            # Optionally, raise an error or convert, for now, let's default to scriptdir
            # or handle as an error appropriately depending on desired behavior.
            # For robustness, let's ensure it's always a string or default if invalid.
            folder_path = scriptdir
        self.config["last_used_folder"] = folder_path
        log.debug(f"Set last_used_folder to: {folder_path}")
        # Note: Saving the config should be handled explicitly by the caller if needed immediately.
        # For example, by calling gui_config.save_config(gui_config.config)