Source code for prsctrl.scripts.compression

"""
Functions to compress the full-data (binary, pickle) file in the final data directories using gzip.
Compression reduces the overall size of a data directory (including plots and csvs, which remain uncomopressed) to ~1/3.
The PrsData class can automatically decompress the data when loading it from the directory, so compression is not inconvenient.
New data can be automatically compressed by passing `compress=True` to the `write_full_binary_data()` method of `PrsData`.
"""
import os.path
import shutil
from ..data.util import add_absorption_data, find_all_data_dirs, load_rta, find_all_rta_data_dirs, get_data_name_template, write_rta_csv
from multiprocessing.pool import Pool
from ..data.prsdata import FULL_DATA_SUFFIX, PrsData
import gzip

import logging
log = logging.getLogger(__name__)




def compress_data(datadir: str, remove_uncompressed: bool=True):
    """
    Compress the full data (pickle) file of the data directory in-place.

    :param datadir: Data directory
    :param remove_uncompressed: If True, the original (uncompressed) file is deleted after compression
    :raises FileNotFoundError: 
    :return: tuple of (datadir, Exception, Traceback) if an error occures or (datadir, None, None) if successful (this make batch-processing with process pools more convenient)
    """
    try:
        data = PrsData(datadir)
        full_data_file = os.path.join(datadir, os.path.basename(datadir) + FULL_DATA_SUFFIX)
        if not os.path.isfile(full_data_file):
            if os.path.isfile(full_data_file + ".gz"):
                log.debug(f"Processing '{full_data_file}' - already compressed")
                return  (datadir, None, None)
            raise FileNotFoundError(f"No full data file found: '{full_data_file}'")
        log.debug(f"Processing '{full_data_file}'")
        with open(full_data_file, "rb") as file:
            data = file.read()
        with gzip.open(full_data_file + ".gz", "wb") as zfile:
            zfile.write(data)
        if remove_uncompressed:
            os.remove(full_data_file)

        return  (datadir, None, None)
    except Exception as e:
        import traceback
        tb = traceback.format_exc()
        return (datadir, e, tb)





def compress_all_datas(base_dir: str, **kw):
    """
    Recursively compress the full data files in all data directories in `base_dir`

    """
    all_dirs = find_all_data_dirs(base_dir)
    all_dirs.sort()

    answer = input("WARNING: This will compress all full data .pkl files in place! - Do you have a backup? (y/n): ")
    if answer not in ["y", "Y", "j", "J", "yes", "Yes"]:
        print(f"Then make a backup first.")
        return

    # 1) RTA datas
    with Pool() as pool:
        errs = pool.starmap(compress_data, [(data_dir, ) for data_dir in all_dirs])
    n_fail = 0
    for (d, e, tb) in errs:
        if e is None: continue
        print(f"\n\nFAILED DATA: {d}")
        print(tb)
        print(e)
        n_fail += 1

    print(f"\n\n{len(errs)} total, {n_fail} failures")