Source code for GribUtil

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#*******************************************************************************
# @Author: Anne Fouilloux (University of Oslo)
#
# @Date: July 2014
#
# @Change History:
#   February 2018 - Anne Philipp (University of Vienna):
#        - applied PEP8 style guide
#        - added documentation
#        - changed some naming
#
# @License:
#    (C) Copyright 2014-2018.
#
#    This software is licensed under the terms of the Apache Licence Version 2.0
#    which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
#
# @Class Description:
#    The GRIB API provides all necessary tools to work directly with the
#    grib files. Nevertheless, the GRIB API tools are very basic and are in
#    direct connection with the grib files. This class provides some higher
#    functions which apply a set of GRIB API tools together in the respective
#    context. So, the class initially contains a list of grib files (their
#    names) and the using program then applies the methods directly on the
#    class objects without having to think about how the actual GRIB API
#    tools have to be arranged.
#
# @Class Content:
#    - __init__
#    - get_keys
#    - set_keys
#    - copy_dummy_msg
#    - index
#
# @Class Attributes:
#    - filenames
#
#*******************************************************************************

# ------------------------------------------------------------------------------
# MODULES
# ------------------------------------------------------------------------------
import os

from eccodes import (codes_grib_new_from_file, codes_is_defined, codes_get,
                     codes_release, codes_set, codes_write, codes_index_read,
                     codes_index_new_from_file, codes_index_add_file,
                     codes_index_write, codes_set_values)

# ------------------------------------------------------------------------------
# CLASS
# ------------------------------------------------------------------------------
[docs]class GribUtil(object): ''' Class for GRIB utilities (new methods) based on GRIB API ''' # -------------------------------------------------------------------------- # CLASS FUNCTIONS # -------------------------------------------------------------------------- def __init__(self, filenames): '''Initialise an object of GribUtil and assign a list of filenames. Parameters ---------- filenames : :obj:`list` of :obj:`strings` A list of filenames. Return ------ ''' self.filenames = filenames return
[docs] def get_keys(self, keynames, wherekeynames=[], wherekeyvalues=[]): '''Get keyvalues for a given list of keynames a where statement can be given (list of key and list of values) Parameters ---------- keynames : :obj:`list` of :obj:`string` List of keynames. wherekeynames : :obj:`list` of :obj:`string`, optional Default value is an empty list. wherekeyvalues : :obj:`list` of :obj:`string`, optional Default value is an empty list. Return ------ return_list : :obj:`list` of :obj:`string` List of keyvalues for given keynames. ''' fileid = open(self.filenames, 'r') return_list = [] while 1: gid = codes_new_from_file(fileid) if gid is None: break if len(wherekeynames) != len(wherekeyvalues): raise Exception("Number of key values and key names must be \ the same. Give a value for each keyname!") select = True i = 0 for wherekey in wherekeynames: if not codes_is_defined(gid, wherekey): raise Exception("where key was not defined") select = (select and (str(wherekeyvalues[i]) == str(codes_get(gid, wherekey)))) i += 1 if select: llist = [] for key in keynames: llist.extend([str(codes_get(gid, key))]) return_list.append(llist) codes_release(gid) fileid.close() return return_list
[docs] def set_keys(self, fromfile, keynames, keyvalues, wherekeynames=[], wherekeyvalues=[], strict=False, filemode='w'): '''Opens the file to read the grib messages and then write the selected messages (with wherekeys) to a new output file. Also, the keyvalues of the passed list of keynames are set. Parameters ---------- fromfile : :obj:`string` Filename of the input file to read the grib messages from. keynames : :obj:`list` of :obj:`string` List of keynames to set in the selected messages. Default is an empty list. keyvalues : :obj:`list` of :obj:`string` List of keyvalues to set in the selected messages. Default is an empty list. wherekeynames : :obj:`list` of :obj:`string`, optional List of keynames to select correct message. Default value is an empty list. wherekeyvalues : :obj:`list` of :obj:`string`, optional List of keyvalues for keynames to select correct message. Default value is an empty list. strict : :obj:`boolean`, optional Decides if everything from keynames and keyvalues is written out the grib file (False) or only those meeting the where statement (True). Default is False. filemode : :obj:`string`, optional Sets the mode for the output file. Default is "w". Return ------ ''' if len(wherekeynames) != len(wherekeyvalues): raise Exception("Give a value for each keyname!") fout = open(self.filenames, filemode) fin = open(fromfile) while 1: gid = codes_grib_new_from_file(fin) if gid is None: break select = True i = 0 for wherekey in wherekeynames: if not codes_is_defined(gid, wherekey): raise Exception("wherekey was not defined") select = (select and (str(wherekeyvalues[i]) == str(codes_get(gid, wherekey)))) i += 1 if select: i = 0 for key in keynames: if key == 'values': codes_set_values(gid, keyvalues[i]) else: codes_set(gid, key, keyvalues[i]) i += 1 codes_write(gid, fout) codes_release(gid) fin.close() fout.close() return
[docs] def copy_dummy_msg(self, filename_in, selectWhere=True, keynames=[], keyvalues=[], filemode='w'): '''Add the content of another input grib file to the objects file but only messages corresponding to keys/values passed to the function. The selectWhere switch decides if to copy the keys equal to (True) or different to (False) the keynames/keyvalues list passed to the function. Parameters ---------- filename_in : :obj:`string` Filename of the input file to read the grib messages from. selectWhere : :obj:`boolean`, optional Decides if to copy the keynames and values equal to (True) or different to (False) the keynames/keyvalues list passed to the function. Default is True. keynames : :obj:`list` of :obj:`string`, optional List of keynames. Default is an empty list. keyvalues : :obj:`list` of :obj:`string`, optional List of keyvalues. Default is an empty list. filemode : :obj:`string`, optional Sets the mode for the output file. Default is "w". Return ------ ''' if len(keynames) != len(keyvalues): raise Exception("Give a value for each keyname!") fin = open(filename_in) fout = open(self.filenames, filemode) fields = 0 while fields < 1: gid = codes_grib_new_from_file(fin) if gid is None: break select = True i = 0 for key in keynames: if not codes_is_defined(gid, key): raise Exception("Key was not defined") if selectWhere: select = (select and (str(keyvalues[i]) == str(codes_get(gid, key)))) else: select = (select and (str(keyvalues[i]) != str(codes_get(gid, key)))) i += 1 if select: fields = fields + 1 codes_write(gid, fout) codes_release(gid) fin.close() fout.close() return
[docs] def index(self, index_keys=["mars"], index_file="my.idx"): '''Create index file from a list of files if it does not exist or read an index file. Parameters ---------- index_keys: :obj:`list` of :obj:`string`, optional Contains the list of key parameter names from which the index is to be created. Default is a list with a single entry string "mars". index_file: :obj:`string`, optional Filename where the indices are stored. Default is "my.idx". Return ------ iid: :obj:`integer` Grib index id. ''' print("... index will be done") iid = None if os.path.exists(index_file): iid = codes_index_read(index_file) print("Use existing index file: %s " % (index_file)) else: for filename in self.filenames: print("Inputfile: %s " % (filename)) if iid is None: iid = codes_index_new_from_file(filename, index_keys) else: codes_index_add_file(iid, filename) if iid is not None: codes_index_write(iid, index_file) print('... index done') return iid