#!/usr/bin/env python

"""read the xml files and create a shell script that

can be sourced to set the environment needed for standalone NCL testing

__________________________

Created on October, 2015

Author: CSEG <cseg@cgd.ucar.edu>

"""

from __future__ import print_function

import sys

# check the system python version and require 2.7.x or greater

if sys.hexversion < 0x02070000:

print(70 * "*")

print("ERROR: {0} requires python >= 2.7.x. ".format(sys.argv[0]))

print("It appears that you are running python {0}".format(

".".join(str(x) for x in sys.version_info[0:3])))

print(70 * "*")

sys.exit(1)

#

# built-in modules

#

import argparse

import os

import re

import stat

try:

import lxml.etree as etree

except:

import xml.etree.ElementTree as etree

# get the postprocess virtualenv path from the env_postprocess.xml file

env_file = './env_postprocess.xml'

postprocess_path = ''

standalone = ''

if os.path.isfile(env_file):

xml_tree = etree.ElementTree()

xml_tree.parse(env_file)

for entry_tag in xml_tree.findall('entry'):

if entry_tag.get('id') == 'POSTPROCESS_PATH':

postprocess_path = entry_tag.get('value')

if entry_tag.get('id') == 'STANDALONE':

standalone = entry_tag.get('value')

else:

err_msg = ('pp_config ERROR: env_postprocess.xml does not exist in this directory.')

raise OSError(err_msg)

# check if virtualenv is activated

if hasattr(sys, 'real_prefix'):

try:

import jinja2

except:

#

# activate the virtual environment that was created by create_python_env.sh

#

activate_file = '{0}/cesm-env2/bin/activate_this.py'.format(postprocess_path)

if not os.path.isfile(activate_file):

err_msg = ('pp_config ERROR: the virtual environment in {0} does not exist.'.format(postprocess_path) \

+ 'Please run {0}/create_python_env.sh -cimeroot [cimeroot] -machine [machine name]'.format(postprocess_path))

raise OSError(err_msg)

try:

execfile(activate_file, dict(__file__=activate_file))

except:

raise OSError('pp_config ERROR: Unable to activate python virtualenv {0}'.format(activate_file))

else:

#

# activate the virtual environment that was created by create_python_env.sh

#

activate_file = '{0}/cesm-env2/bin/activate_this.py'.format(postprocess_path)

if not os.path.isfile(activate_file):

err_msg = ('pp_config ERROR: the virtual environment in {0} does not exist.'.format(postprocess_path) \

+ 'Please run {0}/create_python_env.sh -cimeroot [cimeroot] -machine [machine name]'.format(postprocess_path))

raise OSError(err_msg)

try:

execfile(activate_file, dict(__file__=activate_file))

except:

raise OSError('pp_config ERROR: Unable to activate python virtualenv {0}'.format(activate_file))

if sys.version_info[0] == 2:

from ConfigParser import SafeConfigParser as config_parser

else:

from configparser import ConfigParser as config_parser

re_val = re.compile(r'\$(\{([A-Za-z0-9_]+)\}|[A-Za-z0-9_]+)')

#

# import modules installed in the virtual environment

#

from diag_utils import diagUtilsLib

# -------------------------------------------------------------------------------

# commandline_options - parse any command line options

# -------------------------------------------------------------------------------

def commandline_options():

"""Process the command line arguments.

"""

parser = argparse.ArgumentParser(

description='Read the caseroot postprocessing XML files and set the local environment.')

parser.add_argument('-component', '--component', choices=['atm','ice','lnd','ocn'], required=True, nargs=1,

help='select a component [atm, ice, lnd, ocn] to set the environment using the associated XML file')

parser.add_argument('-shell', '--shell', choices=['bash','csh'], required=True, nargs=1,

help='select a shell [bash, csh] to create a file that can be sourced')

parser.add_argument('-caseroot', '--caseroot', required=False, nargs=1, default='.',

help='fully quailfied path to case root post processing directory. Defaults to current directory.')

parser.add_argument('-backtrace', '--backtrace', action='store_true',

help='show exception backtraces as extra debugging output')

parser.add_argument('-debug', '--debug', nargs=1, required=False, type=int, default=0,

help='debugging verbosity level output: 0 = none, 1 = minimum, 2 = maximum. 0 is default')

options = parser.parse_args()

return options

#==============================================================================

# expand recursive function

#==============================================================================

def expand (val, src):

"""

Recursively traverse the environment dictionary to expand all environment variables

Arguments

val (string) - value to be expanded

src (dictionary) - source dictionary to search and resolve param val

Return

expanded value

"""

return re_val.sub(lambda m: expand(src.get(m.group(1), ''), src), val)

#==============================================================================

# readXML - read and resolve an XML input file against a source directory

#==============================================================================

def readXML(casedir, env_file_list):

"""

for the <entry id=value> xml files in casedir/env_file_list,

returns a dictionary output["id"]="value"

Arguments:

casedir (string) - case root directory

env_file_list - list of env_*.xml files to parse in the casedir

Return:

output (dictionary)

"""

output = dict()

for efile in env_file_list:

env_file = '{0}/{1}'.format(casedir, efile)

if os.path.isfile(env_file):

xml_tree = etree.ElementTree()

xml_tree.parse(env_file)

for entry_tag in xml_tree.findall('entry'):

output[entry_tag.get('id')] = entry_tag.get('value')

else:

err_msg = 'cesmEnvLib.readXML ERROR: {0} does not exist or cannot be read'.format(env_file)

raise TypeError(err_msg)

# expand nested environment variables

for k,v in output.iteritems():

output[k] = expand(v, output)

# remove () in output dictionary values

for k,v in output.iteritems():

output[k] = re.sub('[()]', '', v)

# wrap the value in quotes if it contains spaces or commas

if ',' in output[k] or ' ' in output[k]:

new_value = '"{0}"'.format(output[k])

output[k] = new_value

return output

# -------------------------------------------------------------------------------

# main

# -------------------------------------------------------------------------------

def main(options):

""" main

"""

envDict = dict()

if options.caseroot:

caseroot = options.caseroot[0]

prefix_map = {'atm':'ATMDIAG_',

'ice':'ICEDIAG_',

'lnd':'LNDDIAG_',

'ocn':'OCNDIAG_'}

prefix = prefix_map[options.component[0]]

compEnvFile = 'env_diags_{0}.xml'.format(options.component[0])

envFiles = ['env_postprocess.xml', compEnvFile]

envDict = readXML(caseroot, envFiles)

envDict = diagUtilsLib.strip_prefix(envDict, prefix)

envDict['CASEROOT'] = options.caseroot[0]

filename = 'set_env_{0}'.format(options.component[0])

first = True

if options.shell[0] == 'csh':

filename = '{0}/{1}.csh'.format(envDict['CASEROOT'], filename)

else:

filename = '{0}/{1}.sh'.format(envDict['CASEROOT'], filename)

with open (filename, 'w') as envFile:

for key, value in envDict.iteritems():

tmpvalue = None

if 'PATH' in key:

tmpvalue = os.getenv(key)

if options.shell[0] == 'csh':

if tmpvalue is not None:

line = 'setenv {0} {1}:{2}\n\n'.format(key, value, tmpvalue)

else:

line = 'setenv {0} {1}\n\n'.format(key, value)

if first:

instruct = '*** Be sure to run "source {0}"'.format(filename)

first = False

else:

if tmpvalue is not None:

line = 'export {0}={1}:{2}\n\n'.format(key, value, tmpvalue)

else:

line = 'export {0}={1}\n\n'.format(key, value)

if first:

instruct = '*** Be sure to run "source {0}"'.format(filename)

first = False

envFile.write(line)

st = os.stat(filename)

os.chmod(filename, st.st_mode | stat.S_IEXEC)

return '*** Successfully created {0}\n{1}'.format(filename, instruct)

#===================================

if __name__ == "__main__":

options = commandline_options()

try:

status = main(options)

sys.exit(status)

except Exception as error:

print(str(error))

sys.exit(1)