source: flex_extract.git/source/python/classes/GribUtil.py @ c2c5948

ctbtodev
Last change on this file since c2c5948 was 702cbb7, checked in by Anne Philipp <anne.philipp@…>, 5 years ago

renamed GribTools? to GribUtils?
  • Property mode set to 100644
File size: 9.9 KB
Line 
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3#*******************************************************************************
4# @Author: Anne Fouilloux (University of Oslo)
5#
6# @Date: July 2014
7#
8# @Change History:
9#   February 2018 - Anne Philipp (University of Vienna):
10#        - applied PEP8 style guide
11#        - added documentation
12#        - changed some naming
13#
14# @License:
15#    (C) Copyright 2014-2018.
16#
17#    This software is licensed under the terms of the Apache Licence Version 2.0
18#    which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
19#
20# @Class Description:
21#    The GRIB API provides all necessary tools to work directly with the
22#    grib files. Nevertheless, the GRIB API tools are very basic and are in
23#    direct connection with the grib files. This class provides some higher
24#    functions which apply a set of GRIB API tools together in the respective
25#    context. So, the class initially contains a list of grib files (their
26#    names) and the using program then applies the methods directly on the
27#    class objects without having to think about how the actual GRIB API
28#    tools have to be arranged.
29#
30# @Class Content:
31#    - __init__
32#    - get_keys
33#    - set_keys
34#    - copy
35#    - index
36#
37# @Class Attributes:
38#    - filenames
39#
40#*******************************************************************************
41
42# ------------------------------------------------------------------------------
43# MODULES
44# ------------------------------------------------------------------------------
45import os
46
47from eccodes import (codes_grib_new_from_file, codes_is_defined, codes_get,
48                     codes_release, codes_set, codes_write, codes_index_read,
49                     codes_index_new_from_file, codes_index_add_file,
50                     codes_index_write)
51
52# ------------------------------------------------------------------------------
53# CLASS
54# ------------------------------------------------------------------------------
55class GribTools(object):
56    '''
57    Class for GRIB utilities (new methods) based on GRIB API
58    '''
59    # --------------------------------------------------------------------------
60    # CLASS FUNCTIONS
61    # --------------------------------------------------------------------------
62    def __init__(self, filenames):
63        '''Initialise an object of GribTools and assign a list of filenames.
64
65        Parameters
66        ----------
67        filenames : :obj:`list` of :obj:`strings`
68             A list of filenames.
69
70        Return
71        ------
72
73        '''
74
75        self.filenames = filenames
76
77        return
78
79
80    def get_keys(self, keynames, wherekeynames=[], wherekeyvalues=[]):
81        '''Get keyvalues for a given list of keynames a where statement
82        can be given (list of key and list of values)
83
84        Parameters
85        ----------
86        keynames : :obj:`list` of :obj:`string`
87            List of keynames.
88
89        wherekeynames : :obj:`list` of :obj:`string`, optional
90            Default value is an empty list.
91
92        wherekeyvalues : :obj:`list` of :obj:`string`, optional
93            Default value is an empty list.
94
95        Return
96        ------
97        return_list : :obj:`list` of :obj:`string`
98            List of keyvalues for given keynames.
99        '''
100
101        fileid = open(self.filenames, 'r')
102
103        return_list = []
104
105        while 1:
106            gid = codes_new_from_file(fileid)
107
108            if gid is None:
109                break
110
111            if len(wherekeynames) != len(wherekeyvalues):
112                raise Exception("Number of key values and key names must be \
113                                 the same. Give a value for each keyname!")
114
115            select = True
116            i = 0
117            for wherekey in wherekeynames:
118                if not codes_is_defined(gid, wherekey):
119                    raise Exception("where key was not defined")
120
121                select = (select and (str(wherekeyvalues[i]) ==
122                                      str(codes_get(gid, wherekey))))
123                i += 1
124
125            if select:
126                llist = []
127                for key in keynames:
128                    llist.extend([str(codes_get(gid, key))])
129                return_list.append(llist)
130
131            codes_release(gid)
132
133        fileid.close()
134
135        return return_list
136
137
138    def set_keys(self, fromfile, keynames, keyvalues, wherekeynames=[],
139                 wherekeyvalues=[], strict=False, filemode='w'):
140        '''Opens the file to read the grib messages and then write
141        them to a new output file. By default all messages are
142        written out. Also, the keyvalues of the passed list of
143        keynames are set or only those meeting the where statement.
144        (list of key and list of values).
145
146        Parameters
147        ----------
148        fromfile : :obj:`string`
149            Filename of the input file to read the grib messages from.
150
151        keynames : :obj:`list` of :obj:`string`
152            List of keynames. Default is an empty list.
153
154        keyvalues : :obj:`list` of :obj:`string`
155            List of keynames. Default is an empty list.
156
157        wherekeynames : :obj:`list` of :obj:`string`, optional
158            Default value is an empty list.
159
160        wherekeyvalues : :obj:`list` of :obj:`string`, optional
161            Default value is an empty list.
162
163        strict : :obj:`boolean`, optional
164            Decides if everything from keynames and keyvalues
165            is written out the grib file (False) or only those
166            meeting the where statement (True). Default is False.
167
168        filemode : :obj:`string`, optional
169            Sets the mode for the output file. Default is "w".
170
171        Return
172        ------
173
174        '''
175        fout = open(self.filenames, filemode)
176        fin = open(fromfile)
177
178        while 1:
179            gid = codes_new_from_file(fin)
180
181            if gid is None:
182                break
183
184            if len(wherekeynames) != len(wherekeyvalues):
185                raise Exception("Give a value for each keyname!")
186
187            select = True
188            i = 0
189            for wherekey in wherekeynames:
190                if not codes_is_defined(gid, wherekey):
191                    raise Exception("where Key was not defined")
192
193                select = (select and (str(wherekeyvalues[i]) ==
194                                      str(codes_get(gid, wherekey))))
195                i += 1
196
197            if select:
198                i = 0
199                for key in keynames:
200                    codes_set(gid, key, keyvalues[i])
201                    i += 1
202
203            codes_write(gid, fout)
204
205            codes_release(gid)
206
207        fin.close()
208        fout.close()
209
210        return
211
212    def copy(self, filename_in, selectWhere=True,
213             keynames=[], keyvalues=[], filemode='w'):
214        '''Add the content of another input grib file to the objects file but
215        only messages corresponding to keys/values passed to the function.
216        The selectWhere switch decides if to copy the keys equal to (True) or
217        different to (False) the keynames/keyvalues list passed to the function.
218
219        Parameters
220        ----------
221        filename_in : :obj:`string`
222            Filename of the input file to read the grib messages from.
223
224        selectWhere : :obj:`boolean`, optional
225            Decides if to copy the keynames and values equal to (True) or
226            different to (False) the keynames/keyvalues list passed to the
227            function. Default is True.
228
229        keynames : :obj:`list` of :obj:`string`, optional
230            List of keynames. Default is an empty list.
231
232        keyvalues : :obj:`list` of :obj:`string`, optional
233            List of keynames. Default is an empty list.
234
235        filemode : :obj:`string`, optional
236            Sets the mode for the output file. Default is "w".
237
238        Return
239        ------
240
241        '''
242
243        fin = open(filename_in)
244        fout = open(self.filenames, filemode)
245
246        while 1:
247            gid = codes_new_from_file(fin)
248
249            if gid is None:
250                break
251
252            if len(keynames) != len(keyvalues):
253                raise Exception("Give a value for each keyname!")
254
255            select = True
256            i = 0
257            for key in keynames:
258                if not codes_is_defined(gid, key):
259                    raise Exception("Key was not defined")
260
261                if selectWhere:
262                    select = (select and (str(keyvalues[i]) ==
263                                          str(codes_get(gid, key))))
264                else:
265                    select = (select and (str(keyvalues[i]) !=
266                                          str(codes_get(gid, key))))
267                i += 1
268
269            if select:
270                codes_write(gid, fout)
271
272            codes_release(gid)
273
274        fin.close()
275        fout.close()
276
277        return
278
279    def index(self, index_keys=["mars"], index_file="my.idx"):
280        '''Create index file from a list of files if it does not exist or
281        read an index file.
282
283        Parameters
284        ----------
285        index_keys: :obj:`list` of :obj:`string`, optional
286            Contains the list of key parameter names from
287            which the index is to be created.
288            Default is a list with a single entry string "mars".
289
290        index_file: :obj:`string`, optional
291            Filename where the indices are stored.
292            Default is "my.idx".
293
294        Return
295        ------
296        iid: :obj:`integer`
297            Grib index id.
298        '''
299        print("... index will be done")
300        iid = None
301
302        if os.path.exists(index_file):
303            iid = codes_index_read(index_file)
304            print("Use existing index file: %s " % (index_file))
305        else:
306            for filename in self.filenames:
307                print("Inputfile: %s " % (filename))
308                if iid is None:
309                    iid = codes_index_new_from_file(filename, index_keys)
310                else:
311                    codes_index_add_file(iid, filename)
312
313            if iid is not None:
314                codes_index_write(iid, index_file)
315
316        print('... index done')
317
318        return iid
Note: See TracBrowser for help on using the repository browser.
hosted by ZAMG