desihub / desispec / 10947976957

import os

import glob

import sys

import numpy as np

import re

import time

from astropy.table import Table

from desispec.io.meta import findfile

from desispec.workflow.proctable import get_processing_table_column_defs

from desispec.workflow.utils import define_variable_from_environment, listpath, \

from desispec.workflow.tableio import write_table, load_table

from desispec.scripts.exposuretable import create_exposure_tables

def update_processing_tables(nights=None, night_range=None, orig_filetype='csv',

    if nights is None and night_range is None:

        raise ValueError("Must specify either nights or night_range."

    proctab_template = findfile('proctable', night=99999999)

    proctab_template = proctab_template.replace('99999999', '202[0-9][01][0-9][0-3][0-9]')

    proctab_template = proctab_template.replace('.csv', f'.{orig_filetype}')

    nights_with_proctables = list()

    for ptabfn in glob.glob(proctab_template):

        matches = re.findall('202\d{5}', os.path.basename(ptabfn))

        if len(matches) == 1:

            n = int(matches[0])

            nights_with_proctables.append(n)

            print(f"Couldn't parse a night from proctable file: {ptabfn}")

    check_night = False

    if nights is None or nights == 'all':

        nights = nights_with_proctables

        nights = [int(val.strip()) for val in nights.split(",")]

        check_night = True

    nights = np.sort(nights)

    if night_range is not None:

        if ',' not in night_range:

            raise ValueError("night_range must be a comma separated pair of "

        nightpair = night_range.split(',')

        if len(nightpair) != 2 or not nightpair[0].isnumeric() \

            raise ValueError("night_range must be a comma separated pair of "

        first_night, last_night = nightpair

        nights = nights[np.where(int(first_night) <= nights.astype(int))[0]]

        nights = nights[np.where(int(last_night) >= nights.astype(int))[0]]

    ptab_cols, ptab_dtypes, ptab_defs = get_processing_table_column_defs(return_default_values=True)

    ptab_cols, ptab_dtypes = np.array(ptab_cols), np.array(ptab_dtypes)

    print("Nights: ", nights)

    for night in nights:

        if check_night and night not in nights_with_proctables:

            print(f"Night {night} doesn't have a processing table: Skipping.")

            continue

        orig_pathname = findfile('proctable', night=night).replace('.csv', f'.{orig_filetype}')

        if not os.path.exists(orig_pathname):

            print(f'Could not find processing table for night={night} at:'

            continue

        origtable = load_table(orig_pathname, tabletype='proctab')

        curr_colnames = np.array(list(origtable.colnames))

        expected_cols = np.isin(curr_colnames, ptab_cols)

        found_cols = np.isin(ptab_cols, curr_colnames)

        if np.all(expected_cols) and np.all(found_cols):

            print(f"{orig_pathname} has all of the expected columns, not updating this table.")

            continue

        unexpected = list(curr_colnames[~expected_cols])

        missing = list(ptab_cols[~found_cols])

        print(f"Found the following unexpected columns: {unexpected}")

        print(f"Found the following missing columns: {missing}")

        if 'CAMWORD' in unexpected and 'PROCCAMWORD' in missing:

            print(f"CAMWORD listed instead of PROCCAMWORD. Updating that.")

            origtable.rename_column('CAMWORD', 'PROCCAMWORD')

            unexpected.remove('CAWORD')

            missing.remove('PROCCAMWORD')

        if len(unexpected) > 0:

            print(f"WARNING: Script detected unexpected columns. Only handle "

            for colname in unexpected:

                origtable.remove_column(colname)

        for colname in missing:

            if colname not in ['BADAMPS', 'LASTSTEP', 'EXPFLAG']:

                print(f"WARNING: Script didn't expect {colname} to be missing. "

            colindex = np.where(ptab_cols==colname)[0][0]

            newdat = [ptab_defs[colindex]] * len(origtable)

            newcol = Table.Column(name=colname, data=newdat, dtype=ptab_dtypes[colindex])

            origtable.add_column(newcol)

        origtable = origtable[list(ptab_cols)]

        if dry_run:

            print("\n\nOutput file would have been:")

            origtable.pprint_all()

            ftime = time.strftime("%Y%m%d_%Hh%Mm")

            replaced_pathname = orig_pathname.replace(f".{orig_filetype}",

            print(f"Moving original file from {orig_pathname} to {replaced_pathname}")

            os.rename(orig_pathname,replaced_pathname)

            time.sleep(0.1)

            out_pathname = orig_pathname.replace(f".{orig_filetype}", f".{out_filetype}")

            write_table(origtable, out_pathname)

            print(f"Updated file saved to {out_pathname}. Original archived as {replaced_pathname}")

            print("\n\n")

        sys.stdout.flush()

        sys.stderr.flush()

    print("Exposure table regenerations complete")