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)

    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."
            )

        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