src/ai_dashboard/loader.py

"""This module implements data loading utilities for the AI Dashboard application."""

from typing import Any
import logging
import base64
import io

import pandas as pd

logger = logging.getLogger(__name__)


class DataLoader:
    """Class to handle data loading for the AI Dashboard."""

    def __init__(self, data_path: str):
        """Initialize the DataLoader."""
        logger.debug("DataLoader initialized.")
        self.data_path = data_path

    def load_csv(self) -> Any:
        """Load data from the specified path/url."""
        logger.debug("Loading data from: %s", self.data_path)
        if not self.data_path or self.data_path.strip() == "":
            logger.info("No data_path provided. Waiting for user upload/URL...")
            return None
        try:
            data = pd.read_csv(self.data_path)
            return data
        except Exception as err:  # pylint: disable=W0718
            logger.error("Failed to load data from %s: %s", self.data_path, err)
            raise err

    def load_xlsx(self) -> Any:
        """Load data from the specified path/url."""
        logger.debug("Loading data from: %s", self.data_path)
        try:
            data = pd.read_excel(self.data_path)
            return data
        except Exception as err:  # pylint: disable=W0718
            logger.error("Failed to load data from %s: %s", self.data_path, err)
            raise err

    def load_from_url(self, url: str) -> Any:
        """Load CSV from a direct URL."""
        try:
            df = pd.read_csv(url)
            return df
        except Exception as err:
            logger.error("Failed to load CSV from URL %s: %s", url, err)
            raise err

    def load_uploaded_csv(self, uploaded_contents: str) -> Any:
        """Load CSV from uploaded file contents (Dash Upload)."""
        try:
            _, content_string = uploaded_contents.split(",")
            decoded = base64.b64decode(content_string)
            df = pd.read_csv(io.StringIO(decoded.decode("utf-8")))
            return df
        except Exception as err:
            logger.error("Failed to load uploaded CSV: %s", err)
            raise err