source: tags/libIGCM_v3.0_beta1/libIGCM_post/xios_parser.py

#! /usr/bin/env python
# coding: utf-8

#**************************************************************
# Author: Sebastien Denvil
# Contact: Sebastien.Denvil__at__ipsl.jussieu.fr
# $Revision::                                         $ Revision of last commit
# $Author::                                           $ Author of last commit
# $Date::                                             $ Date of last commit
# IPSL (2006)
#  This software is governed by the CeCILL licence see libIGCM/libIGCM_CeCILL.LIC
#
#**************************************************************

import os, pwd, sys, traceback, argparse
import xml.etree.ElementTree as ET
#import readline, rlcompleter
#readline.parse_and_bind("tab: complete")

indent = 0
currentDepth = 0
ignoreElems = []
fromField=[]
fromFile=[]

def dump(args):
    """Dump XIOS xmls files."""
    # Read and dump xios_def_xml
    for inputFile in args.file:        
        # Read the file_def_xml
        print '\nReading %s \n|' % (inputFile)
        try:
            tree = ET.parse(inputFile)
        except:
            print "Parse error. Please fix so that it can be parsed."
            traceback.print_exc(file=sys.stdout)
            return
        root=tree.getroot()
        # Call the recursive print
        printRecur(root)

def tsquery(args):
    """query timeseries related parameters from an XIOS xml file."""
    if args.verbosity >= 1: print 'Reading timeseries_def_xml=',args.file[0]
    try:
        tree = ET.parse(args.file[0])
    except:
        print "Parse error. Please fix so that it can be parsed."
        traceback.print_exc(file=sys.stdout)
        return
    root=tree.getroot()
    if args.verbosity >= 3: print root.tag, root.attrib
    findTimeSeries(root)
        
def printRecur(root):
    """Recursively prints the tree."""
    global indent
    global currentDepth
    if root.tag in ignoreElems:
        return
    print ' '*indent + '|--> %s: %s' % (root.tag, root.attrib)
    currentDepth += 1
    indent += 4
    if currentDepth <= args.depth or args.depth == None:
        for elem in list(root):
            printRecur(elem)
    currentDepth -= 1
    indent -= 4

def findTimeSeries(root):
    """Recursively find and list field tag with "timeseries", "id", "output_freq" and enable=.TRUE. attribute."""
    if root.tag == 'file' and root.attrib.get('timeseries'):
        if root.attrib.get('enabled') == '.TRUE.':
            print 'id=%s,output_freq=%s' % (root.attrib.get('id'), root.attrib.get('output_freq'))
    else:
        for elem in list(root):
            findTimeSeries(elem)
    
def findField(root):
    """Recursively find and list field tag with "id" or "field_ref" attribute."""
    global fromField
    global fromFile
    if root.tag in ignoreElems:
        return    
    if root.tag == 'field' and root.attrib.get('id'):
        fromField.append(root.attrib.get('id'))
    elif root.tag == 'field' and root.attrib.get('field_ref'):
        fromFile.append(root.attrib.get('field_ref'))
    else:
        for elem in list(root):
            findField(elem)

def findFieldToRemove(root, fieldToRemove):
    """Recursively find tag having a field_ref in fieldToRemove."""
    if args.verbosity >= 3 and root.tag == 'file' and root.attrib.get('id'):
        print '\nFIELDS FROM FILE_DEF with id', root.attrib.get('id')
    for field in root.findall('field'):
        if args.verbosity >= 3: print 'field_ref=', field.attrib.get('field_ref')
        if field.attrib.get('field_ref') in fieldToRemove:
            if args.correction:
                if args.verbosity >= 2: print 'removing ', field.attrib.get('field_ref')
                root.remove(field)
            else:
                if args.verbosity >= 2: print 'To be removed ', field.attrib.get('field_ref')
    for elem in list(root):
        findFieldToRemove(elem, fieldToRemove)
        
def check(args):
    global fromFile
    exitCode=0
    # Read the field_def_xml
    if args.verbosity >= 1: print '\nReading field_def_xml=',args.field[0]
    try:
        tree = ET.parse(args.field[0])
    except:
        print 'Parse error with %s. Please fix so that it can be parsed.' % (args.field[0])
        traceback.print_exc(file=sys.stdout)
        return
    root=tree.getroot()
    if args.verbosity >= 3: print root.tag, root.attrib, '\n'
    # Build a list of field.id from field_def
    findField(root)

    # Loop over file_def files
    for inputFile in args.file:        
        # Read the file_def_xml
        if args.verbosity >= 1: print '\nReading file_def_xml=',inputFile
        try:
            tree = ET.parse(inputFile)
        except:
            print "Parse error. Please fix so that it can be parsed."
            traceback.print_exc(file=sys.stdout)
            return
        root=tree.getroot()
        fromFile=[]
        if args.verbosity >= 3: print root.tag, root.attrib, '\n'

        # Build a list of field_ref from file_def
        findField(root)
        #print '4. fromFile=', fromFile
        
        # Compare the two lists. fromField must be a superset of fromFile.
        if set(fromField).issuperset(set(fromFile)):
            if args.verbosity >= 1: print '\nALL GOOD with %s' % (inputFile)
            if args.verbosity >= 3: print 'fromField=', fromField
            if args.verbosity >= 3: print 'fromFile=', fromFile
        else:
            if args.verbosity >= 1: print '\nTROUBLE AHEAD with %s' % (inputFile)
            if args.verbosity >= 3: print ', '.join(sorted(list(set(fromFile)-set(fromField))))
            # Identify fields in fromFile but not in fromField
            fieldToRemove=list(set(fromFile)-set(fromField))
            if args.verbosity >= 3: print 'fieldToRemove=', fieldToRemove
            # 
            # And now locate and remove them if the modify command has been called
            findFieldToRemove(root, fieldToRemove)
            # Final steps
            if args.correction: tree.write('modified.'+inputFile)
            if not len(fieldToRemove) == 0 and not args.correction:
                exitCode=1
    # The end
    sys.exit(exitCode)

def showtime(args):
    """
    prints table of formatted text format options
    """
    for style in xrange(6):
        for fg in xrange(30,36):
            s1 = ''
            for bg in xrange(40,46):
                format = ';'.join([str(style), str(fg), str(bg)])
                s1 += '\x1b[%sm %s \x1b[0m' % (format, pwd.getpwuid(os.getuid())[4]+' is on fire')
            print s1
        print '\n'
    
if __name__ == '__main__':

    try:
        # Create the top-level parser
        parser = argparse.ArgumentParser(description='XIOS2 xml files tooling and ironsmith')
        subparsers = parser.add_subparsers(description='Dump, check or modify xios xml files')

        # create the parser for the "dump" command
        parser_dump = subparsers.add_parser('dump',help='Dump the xml content without all the xml\'s ironsmith')
        parser_dump.add_argument('-d', '--depth', type=int, default=None, help='How deep do we go. Full tree by default')
        parser_dump.add_argument('file', nargs='+', help='XIOS xml file(s) to dump')
        parser_dump.set_defaults(func=dump)

        # create the parser for the "tsquery" command
        parser_check = subparsers.add_parser('tsquery', help='query timeseries related parameters from an xml file')
        parser_check.add_argument('--file', nargs=1, required=True, help='XIOS xml timeseries_def type')
        parser_check.set_defaults(func=tsquery)
        
        # create the parser for the "check" command
        parser_check = subparsers.add_parser('check', help='Check consistency between field_def and file_def files')
        parser_check.add_argument('--field', nargs=1, required=True, help='XIOS xml field_def type')
        parser_check.add_argument('--file', nargs='+', required=True, help='XIOS xml file_def type')
        parser_check.set_defaults(func=check, correction=False)
        
        # create the parser for the "modify" command
        parser_check = subparsers.add_parser('modify', help='Will make sure field_def is a superset of file_def')
        parser_check.add_argument('--field', nargs=1, required=True, help='XIOS xml field_def type')
        parser_check.add_argument('--file', nargs='+', required=True, help='XIOS xml file_def type')
        parser_check.set_defaults(func=check, correction=True)

        # create the parser for the "modify" command
        parser_check = subparsers.add_parser('showtime', help='Just want to make sure you feel good today')
        parser_check.set_defaults(func=showtime)
        
        # Each possible option
        parser.add_argument('-v', '--verbosity', action='count', default=0)

        # Parse the args.
        args = parser.parse_args()

        # And call whatever function was selected
        args.func(args)
    except KeyboardInterrupt:
        print "Shutdown requested...exiting"
    except Exception:
        traceback.print_exc(file=sys.stdout)
    sys.exit(0)