src/mlpipeline/utils/common.py

import os
import json
import yaml
import pickle
import joblib
from pathlib import Path
from typing import Any, Dict
from box import ConfigBox
from ensure import ensure_annotations

from mlpipeline.logging.logger import get_logger

logger = get_logger(__name__)


@ensure_annotations
def read_yaml(path: Path) -> ConfigBox:
    try:
        with open(path) as yaml_file:
            content = yaml.safe_load(yaml_file)
            logger.info(f"YAML file loaded: {path}")
            return ConfigBox(content)
    except Exception as e:
        logger.error(f"Error reading YAML file {path}: {e}")
        raise e


@ensure_annotations
def create_directories(paths: list, verbose: bool = True):
    for path in paths:
        os.makedirs(path, exist_ok=True)
        if verbose:
            logger.info(f"Created directory: {path}")


@ensure_annotations
def save_json(path: Path, data: Dict):
    with open(path, "w") as f:
        json.dump(data, f, indent=4)
    logger.info(f"JSON file saved: {path}")


@ensure_annotations
def load_json(path: Path) -> ConfigBox:
    with open(path) as f:
        content = json.load(f)
    logger.info(f"JSON file loaded: {path}")
    return ConfigBox(content)


@ensure_annotations
def save_pickle(path: Path, obj: Any):
    with open(path, "wb") as f:
        pickle.dump(obj, f)
    logger.info(f"Pickle file saved: {path}")


@ensure_annotations
def load_pickle(path: Path) -> Any:
    with open(path, "rb") as f:
        obj = pickle.load(f)
    logger.info(f"Pickle file loaded: {path}")
    return obj


@ensure_annotations
def save_model(path: Path, model: Any):
    joblib.dump(model, path)
    logger.info(f"Model saved: {path}")


@ensure_annotations
def load_model(path: Path) -> Any:
    model = joblib.load(path)
    logger.info(f"Model loaded: {path}")
    return model


def save_object(path: Path, obj):
    with open(path, "wb") as f:
        pickle.dump(obj, f)
    logger.info(f"Object saved: {path}")


def load_object(path: Path):
    with open(path, "rb") as f:
        obj = pickle.load(f)
    logger.info(f"Object loaded: {path}")
    return obj


@ensure_annotations
def get_size(path: Path) -> str:
    size_in_kb = round(os.path.getsize(path) / 1024)
    return f"~ {size_in_kb} KB"