Source code for prsctrl.scripts.upgrade_data_version

from os.path import exists
from ..data.upgrade_data import  convert_data
from ..data.util import find_all_data_dirs
from devctrl.utility.file_io import listdir_filter
import os
from multiprocessing.pool import Pool
import logging
log = logging.getLogger(__name__)

import shutil



def convert_data_dirs(base_dir, regex_filter=".*", old_suffix="-original", new_suffix="-converted", move_if_success=True):
    """
    Convert all the data directories in base_dir to the latest prsctrl version format.
    The old data base_dir will be renamed to <base_dir><old_suffix>
    :param move_if_success: If True, the 
    """
    old_base_dir = base_dir.removesuffix("/") + old_suffix
    new_base_dir = base_dir.removesuffix("/") + new_suffix
    os.makedirs(new_base_dir)
    dirs = listdir_filter(base_dir, regex_filter)
    all_success = True
    def f(d):
        if not os.path.isdir(d):
            new_d = d.replace(base_dir, new_base_dir)
            log.debug(f"Renaming {d} -> {new_d}")
            os.rename(d, new_d)
            return (d, "success")
        try:
            convert_data(d, out_dir_base=new_base_dir)
            return (d, "success")
        except Exception as e:
            return (d, e)

    # with ThreadPool(processes=4) as pool:
    #     res = pool.map(f, dirs)
    res = []
    for d in dirs:
        res.append(f(d))

    all_success = True
    for d, e in res:
        print(f"{d:30}: {e}")
        if not e == "success":
            all_success = False
    if all_success and move_if_success:
        print("All conversion succeeded, renaming dirs")
        os.rename(base_dir, old_base_dir)
        os.rename(new_base_dir, base_dir)


def _f(d, base_dir, new_base_dir, old_base_dir):
    new_d = d.replace(base_dir, new_base_dir)
    old_d = d.replace(base_dir, old_base_dir)
    new_base_d = os.path.dirname(new_d)
    old_base_d = os.path.dirname(old_d)
    os.makedirs(new_base_d, exist_ok=True)
    print(d, new_d, old_d, new_base_d, old_base_d)
    try:
        convert_data(d, out_dir_base=new_base_d)
        os.makedirs(old_base_d, exist_ok=True)
        log.info(f"Moving: {d} -> {old_d}")
        log.info(f"Moving: {new_d} -> {d}")
        shutil.move(d, old_d)
        shutil.move(new_d, d)
        return (d, "success")
    except Exception as e:
        return (d, e)




def convert_all_data_dirs(base_dir, old_suffix="-original", new_suffix="-converted", move_if_success=True):
    """
    Convert all the data directories in base_dir to the latest prsctrl version format.
    The old data base_dir will be renamed to <base_dir><old_suffix>
    :param move_if_success: If True, the 
    """
    old_base_dir = base_dir.removesuffix("/") + old_suffix
    new_base_dir = base_dir.removesuffix("/") + new_suffix
    os.makedirs(new_base_dir)
    dirs = find_all_data_dirs(base_dir)
    log.warning("Data dirs " + '\n\t'.join(dirs))
    with Pool(processes=4) as pool:
        res = pool.starmap(_f, [(d, base_dir, new_base_dir, old_base_dir) for d in dirs])
    # res = []
    # for d in dirs:
    #     if d.endswith("-abs"): continue
    #     res.append(f(d))

    all_success = True
    for d, e in res:
        print(f"{d:30}: {e}")
        if not e == "success":
            all_success = False
    if all_success:
        log.info(f"All conversions successful")
    return res