hwrf_pull_inputs.py

#! /usr/bin/env python

import os, sys, logging, getopt, re, StringIO
import produtil.sigsafety, produtil.fileop
import produtil.run, produtil.setup
import hwrf.numerics, hwrf.launcher, hwrf.input

from hwrf.numerics import to_datetime_rel
from produtil.run import exe
from hwrf.input import InputSource, DataCatalog

## @var inputscript
# The script from which the input list is read.
inputscript=None

## @var outputdir
# The output directory where data will reside
outputdir=None

## @var cycles
# The list of YYYYMMDDHH cycles to pull
cycles=None

## @var HOMEhwrf
# Parent directory of HWRF's parm/ and ush/
HOMEhwrf=None

## @var WORKhwrf
# Scrub area, the WORKhwrf variable in the [config] section of conf
WORKhwrf=None

## @var logger
# The logging.Logger object for logging messages.
logger=None

## @var conf_input_list
# A list of input files, relative to parm, to read in when making conf
conf_input_list = [
    'hwrf_input.conf', 'hwrf.conf', 'hwrf_holdvars.conf', 
    'hwrf_basic.conf', 'system.conf' ]

# def setup():
#     """!Sets up the used parts of the produtil package.  This does NOT
#     call produtil.setup.setup().  Instead, it installs the signal
#     handlers and sets up the Python logging module directly.  No other
#     parts of the produtil package are initialized."""
#     # Install SIGTERM, SIGINT, etc. handlers.  Needed to handle batch job
#     # kills correctly.
#     global logger
#     produtil.sigsafety.install_handlers()
#     global handler, logger
#     # Set up logging.  We customize things to look nice.
#     handler=logging.StreamHandler(sys.stderr)
#     handler.setFormatter(logging.Formatter(
#             'pull-inputs:%(levelname)s:%(asctime)s: %(message)s',
#             '%Y%m%d.%H%M%S'))
#     logger=logging.getLogger()
#     logger.setLevel(logging.INFO)
#     logger.addHandler(handler)

def parse_arguments():
    global logger, outputdir, WORKhwrf, HOMEhwrf, cycles, inputscript
    try:
        (optlist,args) = getopt.getopt(sys.argv[1:],'o:w:vi:')
        for opt,val in optlist:
            if opt=='-v':
                logger.setLevel(logging.DEBUG)
                logger.debug('Verbosity enabled.')
            elif opt=='-w':
                logger.info('Work area set to %s'%(val,))
                WORKhwrf=val
            elif opt=='-o':
                logger.info('Output area set to %s'%(val,))
                outputdir=val
            elif opt=='-i':
                logger.info('Input script set to %s'%(val,))
                inputscript=val
            else:
                usage('Invalid option %s'%(opt,))
    except (getopt.GetoptError,ValueError,TypeError) as e:
        usage(str(e))
        sys.exit(1)
    HOMEhwrf=os.environ.get('HOMEhwrf','')
    if not HOMEhwrf:
        HOMEhwrf=os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
    cwd=os.getcwd()
    if not outputdir:
        outputdir=os.path.join(cwd,'hwrfdata')
    if not WORKhwrf:
        WORKhwrf=os.path.join(cwd,'input-temp')
    if inputscript is None:
        if len(args)<2:
            usage('Input script and one cycle must be specified.')
        inputscript=args[0]
        cycles=args[1:]
    elif len(args)<1:
        usage('At least one cycle must be specified.')
    else:
        cycles=args
    logger.info('inputscript is %s'%(inputscript,))
    logger.info('cycles: '+str(cycles))

def make_conf(cycle):
    conf=hwrf.launcher.HWRFLauncher()
    conf.set('config','cycle',cycle)
    conf.cycle=cycle
    assert(conf.cycle is not None)
    for basename in conf_input_list:
        fullname=os.path.join(HOMEhwrf,'parm',basename)
        if not os.path.exists(fullname):
            logger.error('%s: does not exist.  HWRF is not installed in %s'%(
                    fullname,HOMEhwrf))
            sys.exit(1)
        conf.read(fullname)
    conf.set('holdvars','CASE_ROOT','HISTORY')
    conf.set('hwrfdata','inputroot',outputdir)
    conf.set('config','case_root','HISTORY')
    conf.set('config','fcsthist','hist')
    conf.set('config','realtime','false')
    conf.guess_default_values()
    assert(conf.cycle is not None)
    conf.set('dir','HOMEhwrf',HOMEhwrf)
    conf.set('dir','WORKhwrf',WORKhwrf)
    conf.set('dir','com',WORKhwrf)
    return conf

def expand_lists(inputs,stack,names,lists):
    if not names:
        # We have finished iterating.
        inputs.append(stack[-1])
        return
    # Iterate into the next list, expanding all items.
    for item in lists[0]:
        addme=dict(stack[-1])
        addme[names[0]]=item
        stack.append(addme)
        expand_lists(inputs,stack,names[1:],lists[1:])

def pull_command(inputs,dataset,atime,lists,args):
    assert(isinstance(lists,dict))
    pull=dict()
    listflags=set()
    ftimes=list()
    pull.update(dataset=dataset,atime=atime,item=args[0])
    for arg in args[1:]:
        if arg=='optional':
            pull['optional']=True
        elif len(arg)>1 and arg[0]=='+':
            if arg[1:] not in lists:
                raise Exception('%s: list not in %s'%(arg[1:],lists))
            listflags.add(arg[1:])
        elif '0123456789'.find(arg[0])>=0:
            parts=arg.split('..')
            if len(parts)==1:
                fhr=int(parts[0])
                sec=fhr*3600
                ftimes.append(to_datetime_rel(sec ,atime))
            elif len(parts)==2:
                start_fhr=int(parts[0])
                end_fhr=int(parts[1])
                for hr in xrange(start_fhr,end_fhr+1):
                    sec=hr*3600
                    ftimes.append(to_datetime_rel(sec,atime))
            elif len(parts)>2:
                start_fhr=int(parts[0])
                end_fhr=int(parts[2])
                step_fhr=int(parts[1])
                for hr in xrange(start_fhr,end_fhr+1,step_fhr):
                    sec=hr*3600
                    ftimes.append(to_datetime_rel(sec,atime))
        else:
            raise Exception('Unrecognized argument %s'%(arg,))
    if not ftimes:
        ftimes.append(atime)
    listnames = [ listname for listname in listflags ]
    listlists = [ lists[listname] for listname in listnames ]
    for ftime in ftimes:
        addme=dict(pull)
        addme.update(ftime=ftime)
        stack=list()
        stack.append(addme)
        expand_lists(inputs,stack,listnames,listlists)

def make_input_list(conf,inputstream):
    dataset=None
    cycle=conf.cycle
    atime=cycle
    lists=dict()
    inputs=list()
    for line in inputstream:
        stripped=line.strip()
        if stripped[0:1]=='#': continue
        parts=stripped.split()
        if len(parts)<1: continue
        cmd=parts[0].upper()
        if cmd=='DATASET' and len(parts)>1:
            #print 'DATASET',parts[1:]
            dataset=parts[1]
        elif cmd=='ATIME' and len(parts)==2:
            rel=int(parts[1])
            atime=to_datetime_rel(rel*3600,cycle)
        elif cmd=='ATIME' and len(parts)==1:
            atime=cycle
        elif ( cmd=='FILL' or cmd=='FILLINT') and len(parts)==4:
            nums=parts[3].split('..')
            listname=parts[1]
            format=parts[2]
            if len(nums)==1:
                lists[listname]=(format%int(nums[0]))
            elif len(nums)==2:
                lists[listname] = \
                    [str(format)%int(num) \
                         for num in xrange(int(nums[0]),int(nums[1])+1)]
            elif len(nums)>=3:
                lists[listname] = \
                    [str(format)%int(num) \
                      for num in xrange(int(nums[0]),int(nums[2])+1,int(nums[1]))]
            if cmd=='FILLINT':
                old=lists[listname]
                lists[listname]=[ int(O) for O in old ]
        elif cmd=='SET' and len(parts)>2:
            #print 'SET',parts[1:]
            lists[parts[1]]=parts[2:]
        elif cmd=='SETINT' and len(parts)>2:
            #print 'SET',parts[1:]
            lists[parts[1]]=[ int(part) for part in parts[2:] ]
        elif cmd=='PULL' and len(parts)>1:
            #print 'PULL',parts[1:]
            pull_command(inputs,dataset,atime,lists,parts[1:])
        else:
            pass #print 'IGNORE',cmd,parts[1:]
    return inputs

def dump_list(inputs):
    logger.info('List of inputs to pull:')
    for i in inputs:
        s=' '.join([ '%s=%s'%(k,v) for k,v in i.iteritems() ])
        logger.info(s)

def transfer_inputs(inputs,conf,cycle):
    hwrfdata=DataCatalog(conf,'hwrfdata',cycle)
    htar=exe(conf.getexe('htar'))
    hsi=exe(conf.getexe('hsi'))
    insec=conf.getstr('config','input_sources')
    ins=InputSource(conf,insec,cycle,htar,logger,hsi)
    return bool(ins.get(inputs,hwrfdata,False,logger,True))

def main():
    global logger
    produtil.setup.setup(send_dbn=False,thread_logger=True)
    logger=logging.getLogger('pull_inputs')
    parse_arguments()
    okay=True
    for cycle in cycles:
        conf=make_conf(cycle)
        with open(inputscript,'rt') as inf:
            inputs=make_input_list(conf,inf)
        dump_list(inputs)
        okay=okay and transfer_inputs(inputs,conf,cycle)
    if not okay:
        sys.exit(1)

if __name__=='__main__':
    main()