Source code for prsctrl.scripts.edit_metadata

import os
import pickle
import logging

log = logging.getLogger(__name__)

from ..data.prsdata import FULL_DATA_SUFFIX, PrsData, METADATA_FILENAME, TRANSMISSION, REFLECTION
from devctrl.utility.file_io import listdir_filter
from ..data.plot import plot_data



def search_replace_all_in_dir(dir, search, replace, rename_files=False):
    """
    Recursively search through dir.
    Change <search> for <replace> in all csvs (only in comment lines), metadata.pkl and full-data.pkl files.
    # TODO full-data files
    :param dir:
    :param search:
    :param replace:
    :return:
    """
    if not os.path.isdir(dir): return
    log.debug(f"Searching {dir}")
    if rename_files:
        contents = os.listdir(dir)
        for c in contents:
            if search in c:
                new_name = os.path.join(dir, c.replace(search, replace))
                log.info(f"Renaming {c} to {new_name}")
                if os.path.exists(new_name):
                    log.error(f"Can not rename {c} to {new_name}, it exists")
                else:
                    os.rename(os.path.join(dir, c), new_name)
    contents = [os.path.join(dir, c) for c in os.listdir(dir)]
    contents.sort()
    for c in contents:
        if c.endswith(METADATA_FILENAME):
            try:
                with open(c, "rb") as file:
                    md = pickle.load(file)
                changed = False
                for k, v in md.items():
                    if type(v) == str and search in v:
                        md[k] = v.replace(search, replace)
                        changed = True
                if changed:
                    log.info(f"Changes in {c}")
                    with open(c, "wb") as file:
                        pickle.dump(md, file)
            except Exception as e:
                log.error(f"{c}: {e}")
        elif c.endswith(FULL_DATA_SUFFIX):
            try:
                with open(c, "rb") as file:
                    data, md = pickle.load(file)
                changed = False
                for k, v in md.items():
                    if type(v) == str and search in v:
                        md[k] = v.replace(search, replace)
                        changed = True
                if changed:
                    log.info(f"Changes in {c}")
                    with open(c, "wb") as file:
                        pickle.dump((data, md), file)
            except Exception as e:
                log.error(f"{c}: {e}")
        elif c.endswith(".csv"):
            try:
                with open(c, "r") as file:
                    lines = file.readlines()
                changed = False
                for i in range(len(lines)):
                    # only touch comments
                    if not lines[i].startswith("#"): continue
                    if search in lines[i]:
                        changed = True
                        lines[i] = lines[i].replace(search, replace)
                if changed:
                    log.info(f"Changes in {c}")
                    with open(c, "w") as file:
                        file.writelines(lines)
            except Exception as e:
                log.error(f"{c}: {e}")
        else:
            search_replace_all_in_dir(c, search, replace, rename_files=rename_files)





def fix_full_file_metadata(dir):
    """
    Make sure the metadata from the partial file is the same as the full file
    Change <search> for <replace> in all csvs (only in comment lines)
    and in all metadata.pkl files.
    :param dir:
    :return:
    """
    if not os.path.isdir(dir): return
    log.debug(f"Searching {dir}")
    files = os.listdir(dir)
    if not (FULL_DATA_SUFFIX in files and METADATA_FILENAME in files):
        files.sort()
        for f in files:
            fix_full_file_metadata(os.path.join(dir, f))
    else:
        full = os.path.join(dir, FULL_DATA_SUFFIX)
        meta = os.path.join(dir, METADATA_FILENAME)
        # yes we could just copy the metadata from the md file into the full file, but we want to know where sth changed
        with open(full, "rb") as file:
            data, md_full = pickle.load(file)
        with open(meta, "rb") as file:
            md = pickle.load(file)
        changes = False
        for k, v in md.items():
            if k not in md_full:
                print(f"{dir}: key {k} missing in full-data metadata")
                md_full[k] = v
                changes = True
            elif str(md_full[k]) != str(md[k]):
                print(f"{dir}: values for '{k}' differ: md={md[k]}, md_full={md_full[k]}")
                md_full[k] = v
                changes = True
        if changes:
            with open(full, "wb") as file:
                pickle.dump((data, md_full), file)




def fix_missing_ref_metadata(base_dir):
    """
    Fix missing reflectance metadata
    """
    for d in listdir_filter(base_dir, ".*-ref$"):
        data = PrsData(d, file_mode="rw")
        if not data.mode == TRANSMISSION: continue
        print(f"{base_dir}/{d}")
        data.mode = REFLECTION
        data.metadata["mode"] = REFLECTION
        data.metadata["device_lock-in"] = "SR830"
        data.metadata["messages"] = ["Reflectance metadata got deleted :("]
        data.write_metadata()
        data.write_full_file()
        data.save_csv()
        plot_data(data, os.path.basename(data.dirpath), data.dirpath, error_shade=True)