utils/compose_generator.py

"""Utility module for generating Docker Compose files for per-paper compression environments."""

import os
from pathlib import Path
from typing import Optional

import yaml


def generate_paper_compose(paper_name: str, output_dir: str = "outputs") -> str:
    """
    Generate a Docker Compose YAML string for a given paper.

    Args:
        paper_name: Name of the paper (used for output directory)
        output_dir: Base output directory (default: "outputs")

    Returns:
        YAML string representing the Docker Compose configuration
    """
    # Construct the absolute path to the paper's output directory
    abs_output_path = os.path.abspath(os.path.join(output_dir, paper_name))

    # Build the Docker Compose configuration
    compose_config = {
        "version": "3.8",
        "services": {
            "compression-profiler": {
                "image": "compression-profiler:latest",
                "volumes": [f"{abs_output_path}:/work"],
                "working_dir": "/work",
            }
        },
    }

    # Convert to YAML string
    yaml_string = yaml.dump(compose_config, default_flow_style=False)
    return yaml_string


def save_paper_compose(paper_name: str, output_dir: str = "outputs") -> Path:
    """
    Save a Docker Compose file for a given paper.

    Args:
        paper_name: Name of the paper (used for output directory)
        output_dir: Base output directory (default: "outputs")

    Returns:
        Path object pointing to the saved compose file
    """
    # Generate the YAML content
    yaml_content = generate_paper_compose(paper_name, output_dir)

    # Construct the output file path
    abs_output_path = os.path.abspath(os.path.join(output_dir, paper_name))
    compose_file_path = Path(abs_output_path) / ".compose.yaml"

    # Ensure the output directory exists
    compose_file_path.parent.mkdir(parents=True, exist_ok=True)

    # Write the YAML content to file
    with open(compose_file_path, "w") as f:
        f.write(yaml_content)

    return compose_file_path