utils/handle_files.py

import gemmi
import os
import shutil
import gradio as gr
import subprocess
import json
import yaml
from pathlib import Path

def mcif_gz_to_pdb(file_path: str) -> str:
    """
    Converts a .mcif.gz file to pdb and saves it to the same directory. Returns the path to the pdb file.
    
    Parameters:
    ----------
    file_path: str,
        Path to the .mcif.gz file.
    
    Returns
    -------
        str: path to the generated pdb file.
    """
    st = gemmi.read_structure(file_path)
    st.setup_entities()  # Recommended for consistent entity handling [web:18]
    pdb_path = file_path.replace(".cif.gz", ".pdb")
    st.write_minimal_pdb(pdb_path)
    return pdb_path


def download_results_as_zip(directory):
    """
    Check that an output directory is specified, then creates a zip file of the directory for download.
    
    Parameters:
    ----------
    directory: gr.State or str
        Path to the directory containing generated results. None if generation has not been run yet.
    
    Returns
    -------
        str or None: Path to the created zip file if directory is valid, else None.
    """
    if directory is None:
        return None
    zip_path = f"{directory}.zip"
    shutil.make_archive(directory, 'zip', directory)
    return zip_path


def collect_outputs(gen_directory, num_batches, num_designs_per_batch):
    try:
        cmd = f"ls -R {gen_directory}"
        file_list = subprocess.check_output(cmd, shell=True).decode()
        return file_list
    except Exception as e:
        return f"Error: {str(e)}"
    

def load_config(file_path: str | Path) -> dict | list:
    """
    Load YAML or JSON file into a Python object.
    
    Args:
        file_path: Path to the YAML or JSON file.
    
    Returns:
        Parsed Python object (dict, list, etc.).
    
    Raises:
        ValueError: If extension is not .yaml, .yml, or .json.
        Exception: On parse errors.
    """
    path = Path(file_path)
    if not path.exists():
        raise FileNotFoundError(f"File not found: {file_path}")
    
    ext = path.suffix.lower()
    if ext in {'.yaml', '.yml'}:
        with open(path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)  # Secure loader [web:1][web:4]
    elif ext == '.json':
        with open(path, 'r', encoding='utf-8') as f:
            return json.load(f)  # Built-in JSON loader [web:12]
    else:
        raise ValueError(f"Unsupported extension: {ext}. Use .yaml, .yml, or .json.")