diff --git a/data/grib_data_files.txt b/data/grib_data_files.txt index 1a0cb7d99..d25ae3e03 100644 --- a/data/grib_data_files.txt +++ b/data/grib_data_files.txt @@ -22,6 +22,7 @@ index.grib grid_ieee.grib jpeg.grib2 latlon.grib +high_level_api.grib2 lfpw.grib1 missing_field.grib1 missing.grib2 diff --git a/examples/python/CMakeLists.txt b/examples/python/CMakeLists.txt index 7cb1f2798..3dc836d2f 100644 --- a/examples/python/CMakeLists.txt +++ b/examples/python/CMakeLists.txt @@ -60,6 +60,14 @@ list( APPEND tests gts_get_keys metar_get_keys ) + +# The high level python test requires new features in the unittest +# which are only there for python 2.7 onwards +if( PYTHON_VERSION_STRING VERSION_GREATER "2.7" ) + ecbuild_info("Python examples: Adding test for PythonicGrib") + list( APPEND tests high_level_api ) +endif() + foreach( test ${tests} ) ecbuild_add_test( TARGET eccodes_p_${test}_test TYPE SCRIPT diff --git a/examples/python/high_level_api.py b/examples/python/high_level_api.py new file mode 100644 index 000000000..3c83b30e4 --- /dev/null +++ b/examples/python/high_level_api.py @@ -0,0 +1,153 @@ +#!/bin/env python + +""" +Unit tests for ``PythonicGrib``. + +Author: Daniel Lee, DWD, 2014 +""" + +import os +from tempfile import NamedTemporaryFile +import unittest + +from eccodes import GribFile +from eccodes import GribIndex +from eccodes import GribMessage +from eccodes.high_level.gribmessage import IndexNotSelectedError + + +TESTGRIB = "../../data/high_level_api.grib2" +TEST_OUTPUT = "test-output.grib" +TEST_INDEX = "test.index" +TEST_KEYS = ("dataDate", "stepRange") +TEST_VALUES = 20110225, 0 +SELECTION_DICTIONARY = {} +for i1 in range(len(TEST_KEYS)): + SELECTION_DICTIONARY[TEST_KEYS[i1]] = TEST_VALUES[i1] +TEST_INDEX_OUTPUT = TESTGRIB +TEST_STEPRANGE = ('0', '12', '18', '24', '6') + + +class TestGribFile(unittest.TestCase): + """Test GribFile functionality.""" + + def test_memory_management(self): + """Messages in GribFile can be opened and closed properly.""" + with GribFile(TESTGRIB) as grib: + self.assertEqual(len(grib), 5) + for i in range(len(grib)): + msg = GribMessage(grib) + self.assertEqual(msg["shortName"], "msl") + self.assertEqual(len(grib.open_messages), 5) + self.assertEqual(len(grib.open_messages), 0) + + def test_iteration_works(self): + """The GribFile allows proper iteration over all messages.""" + step_ranges = [] + with GribFile(TESTGRIB) as grib: + for _ in range(len(grib)): + msg = GribMessage(grib) + step_ranges.append(msg["stepRange"]) + self.assertSequenceEqual(step_ranges, ["0", "6", "12", "18", "24"]) + + def test_iterator_protocol(self): + """The GribFile allows pythonic iteration over all messages.""" + step_ranges = [] + with GribFile(TESTGRIB) as grib: + for msg in grib: + step_ranges.append(msg["stepRange"]) + self.assertSequenceEqual(step_ranges, ["0", "6", "12", "18", "24"]) + + def test_read_past_last_message(self): + """Trying to open message on exhausted GRIB file raises IOError.""" + with GribFile(TESTGRIB) as grib: + for _ in range(len(grib)): + GribMessage(grib) + self.assertRaises(IOError, lambda: GribMessage(grib)) + + def test_read_invalid_file(self): + """Trying to open message on nonexistent GRIB file raises IOError.""" + with NamedTemporaryFile(mode='r') as f: + with GribFile(f.name) as grib: + self.assertRaises(IOError, lambda: GribMessage(grib)) + + +class TestGribMessage(unittest.TestCase): + """Test GribMessage functionality.""" + + def test_metadata(self): + """Metadata is read correctly from GribMessage.""" + with GribFile(TESTGRIB) as grib: + msg = GribMessage(grib) + key_count = 251 + self.assertEqual(len(msg), key_count) + self.assertEqual(msg.size(), 160219) + self.assertEqual(len(msg.keys()), key_count) + + def test_missing_message_behavior(self): + """Missing messages are detected properly.""" + with GribFile(TESTGRIB) as grib: + msg = GribMessage(grib) + self.assertTrue(msg.missing("scaleFactorOfSecondFixedSurface")) + msg["scaleFactorOfSecondFixedSurface"] = 5 + msg.set_missing("scaleFactorOfSecondFixedSurface") + with self.assertRaises(KeyError): + msg["scaleFactorOfSecondFixedSurface"] + + def test_value_setting(self): + """Keys can be set properly.""" + with GribFile(TESTGRIB) as grib: + msg = GribMessage(grib) + msg["scaleFactorOfSecondFixedSurface"] = 5 + msg["values"] = [1, 2, 3] + + def test_serialize(self): + """Message can be serialized to file.""" + with GribFile(TESTGRIB) as grib: + msg = GribMessage(grib) + with open(TEST_OUTPUT, "w") as test: + msg.write(test) + os.unlink(TEST_OUTPUT) + + def test_clone(self): + """Messages can be used to produce clone Messages.""" + with GribFile(TESTGRIB) as grib: + msg = GribMessage(grib) + msg2 = GribMessage(clone=msg) + self.assertSequenceEqual(msg.keys(), msg2.keys()) + + +class TestGribIndex(unittest.TestCase): + """Test GribIndex functionality.""" + + def test_memory_management(self): + """GribIndex closes GribMessages properly.""" + with GribIndex(TESTGRIB, TEST_KEYS) as idx: + idx.select(SELECTION_DICTIONARY) + self.assertEqual(len(idx.open_messages), 1) + self.assertEqual(len(idx.open_messages), 0) + + def test_create_and_serialize_index(self): + """GribIndex can be saved to file, file can be added to index.""" + with GribIndex(TESTGRIB, TEST_KEYS) as idx: + idx.write(TEST_INDEX) + with GribIndex(file_index=TEST_INDEX) as idx: + idx.add(TESTGRIB) + os.unlink(TEST_INDEX) + + # TODO: Following test disabled due to error message: + # GRIB_API ERROR: please select a value for index key "dataDate" + # Must investigate + # + def _test_index_comprehension(self): + """GribIndex understands underlying GRIB index properly.""" + with GribIndex(TESTGRIB, TEST_KEYS) as idx: + self.assertEqual(idx.size(TEST_KEYS[1]), 5) + self.assertSequenceEqual(idx.values(TEST_KEYS[1]), TEST_STEPRANGE) + with self.assertRaises(IndexNotSelectedError): + idx.select({TEST_KEYS[1]: TEST_VALUES[0]}) + idx.select(SELECTION_DICTIONARY) + self.assertEqual(len(idx.open_messages), 1) + +if __name__ == "__main__": + unittest.main() diff --git a/examples/python/high_level_api.sh b/examples/python/high_level_api.sh new file mode 100755 index 000000000..6c468eaaf --- /dev/null +++ b/examples/python/high_level_api.sh @@ -0,0 +1,10 @@ +#!/bin/sh + +. ./include.sh + +# To get verbose output +#$PYTHON -m unittest -v high_level_api + +$PYTHON $examples_src/high_level_api.py + +rm -f test.index diff --git a/python/eccodes/__init__.py b/python/eccodes/__init__.py index b6c4997b7..4a0f5fd11 100644 --- a/python/eccodes/__init__.py +++ b/python/eccodes/__init__.py @@ -1,2 +1,8 @@ +from __future__ import absolute_import + from .eccodes import * from .eccodes import __version__ + +from .high_level.gribfile import GribFile +from .high_level.gribmessage import GribMessage +from .high_level.gribindex import GribIndex diff --git a/python/eccodes/high_level/__init__.py b/python/eccodes/high_level/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/eccodes/high_level/gribfile.py b/python/eccodes/high_level/gribfile.py new file mode 100644 index 000000000..6df183117 --- /dev/null +++ b/python/eccodes/high_level/gribfile.py @@ -0,0 +1,65 @@ +""" +``GribFile`` class that implements a GRIB file that closes itself and its +messages when it is no longer needed. + +Author: Daniel Lee, DWD, 2014 +""" + +from .. import eccodes +from .gribmessage import GribMessage + + +class GribFile(file): + """ + A GRIB file handle meant for use in a context manager. + + Individual messages can be accessed using the ``next`` method. Of course, + it is also possible to iterate over each message in the file:: + + >>> with GribFile(filename) as grib: + ... # Print number of messages in file + ... len(grib) + ... # Open all messages in file + ... for msg in grib: + ... print(msg["shortName"]) + ... len(grib.open_messages) + >>> # When the file is closed, any open messages are closed + >>> len(grib.open_messages) + """ + + def __enter__(self): + return self + + def __exit__(self, type, value, traceback): + """Close all open messages, release GRIB file handle and close file.""" + while self.open_messages: + self.open_messages.pop().close() + self.file_handle.close() + + def close(self): + """Possibility to manually close file.""" + self.__exit__(None, None, None) + + def __len__(self): + """Return total messages in GRIB file.""" + return eccodes.codes_count_in_file(self.file_handle) + + def __iter__(self): + return self + + def next(self): + try: + return GribMessage(self) + except IOError: + raise StopIteration() + + def __init__(self, filename, mode="r"): + """Open file and receive GRIB file handle.""" + #: File handle for working with actual file on disc + #: The class holds the file it works with because the GRIB API's + #: typechecking does not allow using inherited classes. + self.file_handle = open(filename, mode) + #: Number of message in GRIB file currently being read + self.message = 0 + #: Open messages + self.open_messages = [] diff --git a/python/eccodes/high_level/gribindex.py b/python/eccodes/high_level/gribindex.py new file mode 100644 index 000000000..b0bfabe1d --- /dev/null +++ b/python/eccodes/high_level/gribindex.py @@ -0,0 +1,102 @@ +""" +``GribIndex`` class that implements a GRIB index that allows access to +the GRIB API's index functionality. + +Author: Daniel Lee, DWD, 2014 +""" + +from .. import eccodes +from .gribmessage import GribMessage + + +class GribIndex(object): + """ + A GRIB index meant for use in a context manager. + + Usage:: + + >>> # Create index from file with keys + >>> with GribIndex(filename, keys) as idx: + ... # Write index to file + ... idx.write(index_file) + >>> # Read index from file + >>> with GribIndex(file_index=index_file) as idx: + ... # Add new file to index + ... idx.add(other_filename) + ... # Report number of unique values for given key + ... idx.size(key) + ... # Report unique values indexed by key + ... idx.values(key) + ... # Request GribMessage matching key, value + ... msg = idx.select({key: value}) + """ + + def __enter__(self): + return self + + def __exit__(self, type, value, traceback): + """Release GRIB message handle and inform file of release.""" + while self.open_messages: + self.open_messages[0].close() + eccodes.codes_index_release(self.iid) + + def close(self): + """Possibility to manually close index.""" + self.__exit__(None, None, None) + + def __init__(self, filename=None, keys=None, file_index=None, + grib_index=None): + """ + Create new GRIB index over ``keys`` from ``filename``. + + ``filename`` should be a string of the desired file's filename. + ``keys`` should be a sequence of keys to index. ``file_index`` should + be a string of the file that the index should be loaded from. + ``grib_index`` should be another ``GribIndex``. + + If ``filename`` and ``keys`` are provided, the ``GribIndex`` is + initialized over the given keys from the given file. If they are not + provided, the ``GribIndex`` is read from ``indexfile``. If + ``grib_index`` is provided, it is cloned from the given ``GribIndex``. + """ + #: Grib index ID + self.iid = None + if filename and keys: + self.iid = eccodes.codes_index_new_from_file(filename, keys) + elif file_index: + self.iid = eccodes.codes_index_read(file_index) + elif grib_index: + self.iid = eccodes.codes_new_from_index(grib_index.iid) + else: + raise RuntimeError("No source was supplied " + "(possibilities: grib_file, clone, sample).") + #: Indexed keys. Only available if GRIB is initialized from file. + self.keys = keys + #: Open GRIB messages + self.open_messages = [] + + def size(self, key): + """Return number of distinct values for index key.""" + return eccodes.codes_index_get_size(self.iid, key) + + def values(self, key, type=str): + """Return distinct values of index key.""" + return eccodes.codes_index_get(self.iid, key, type) + + def add(self, filename): + """Add ``filename`` to the ``GribIndex``.""" + eccodes.codes_index_add_file(self.iid, filename) + + def write(self, outfile): + """Write index to filename at ``outfile``.""" + eccodes.codes_index_write(self.iid, outfile) + + def select(self, key_value_pairs): + """ + Return message associated with given key value pairs. + + ``key_value_pairs`` should be passed as a dictionary. + """ + for key in key_value_pairs: + eccodes.codes_index_select(self.iid, key, key_value_pairs[key]) + return GribMessage(gribindex=self) diff --git a/python/eccodes/high_level/gribmessage.py b/python/eccodes/high_level/gribmessage.py new file mode 100644 index 000000000..5c5976521 --- /dev/null +++ b/python/eccodes/high_level/gribmessage.py @@ -0,0 +1,198 @@ +""" +``GribMessage`` class that implements a GRIB message that allows access to +the message's key-value pairs in a dictionary-like manner and closes the +message when it is no longer needed, coordinating this with its host file. + +Author: Daniel Lee, DWD, 2014 +""" + +import collections + +from .. import eccodes + + +class IndexNotSelectedError(Exception): + """GRIB index was requested before selecting key/value pairs.""" + + +class GribMessage(object): + """ + A GRIB message. + + Each ``GribMessage`` is stored as a key/value pair in a dictionary-like + structure. It can be used in a context manager or by itself. When the + ``GribFile`` it belongs to is closed, it closes any open ``GribMessage``s + that belong to it. If a ``GribMessage`` is closed before its ``GribFile`` + is closed, it informs the ``GribFile`` of its closure. + + Scalar and vector values are set appropriately through the same method. + + Usage:: + + >>> with GribFile(filename) as grib: + ... # Access shortNames of all messages + ... for msg in grib: + ... print(msg["shortName"]) + ... # Report number of keys in message + ... len(msg) + ... # Report message size in bytes + ... msg.size + ... # Report keys in message + ... msg.keys() + ... # Check if value is missing + ... msg.missing("scaleFactorOfSecondFixedSurface") + ... # Set scalar value + ... msg["scaleFactorOfSecondFixedSurface"] = 5 + ... # Check key's value + ... msg["scaleFactorOfSecondFixedSurface"] + ... # Set value to missing + ... msg.set_missing("scaleFactorOfSecondFixedSurface") + ... # Missing values raise exception when read with dict notation + ... msg["scaleFactorOfSecondFixedSurface"] + ... # Array values are set transparently + ... msg["values"] = [1, 2, 3] + ... # Messages can be written to file + ... with open(testfile, "w") as test: + ... msg.write(test) + ... # Messages can be cloned from other messages + ... msg2 = GribMessage(clone=msg) + ... # If desired, messages can be closed manually or used in with + ... msg.close() + """ + def __enter__(self): + return self + + def __exit__(self, exception_type, exception_value, traceback): + """Release GRIB message handle and inform file of release.""" + # This assert should never trigger + # assert self.gid is not None + eccodes.codes_release(self.gid) + if self.grib_index: + self.grib_index.open_messages.remove(self) + + def close(self): + """Possibility to manually close message.""" + self.__exit__(None, None, None) + + def __contains__(self, key): + """Check whether a key is present in message.""" + return key in self.keys() + + def __len__(self): + """Return key count.""" + return len(self.keys()) + + def __getitem__(self, key): + """Return value associated with key as its native type.""" + return self.get(key) + + def __setitem__(self, key, value): + """ + Set value associated with key. + + If the object is iterable, + """ + # Alternative implemented (TODO: evaluate) + # if eccodes.codes_get_size(self.gid, key) > 1: + # eccodes.codes_set_array(self.gid, key, value) + # else: + # eccodes.codes_set(self.gid, key, value) + + # Passed value is iterable and not string + if (isinstance(value, collections.Iterable) and not + isinstance(value, basestring)): + eccodes.codes_set_array(self.gid, key, value) + else: + eccodes.codes_set(self.gid, key, value) + + def __iter__(self): + return iter(self.keys()) + + # Not yet implemented + # def itervalues(self): + # return self.values() + + def items(self): + """Return list of tuples of all key/value pairs.""" + return [(key, self[key]) for key in self.keys()] + + def keys(self, namespace=None): + """Get available keys in message.""" + iterator = eccodes.codes_keys_iterator_new(self.gid, namespace=namespace) + keys = [] + while eccodes.codes_keys_iterator_next(iterator): + key = eccodes.codes_keys_iterator_get_name(iterator) + keys.append(key) + eccodes.codes_keys_iterator_delete(iterator) + return keys + + def __init__(self, grib_file=None, clone=None, sample=None, gribindex=None): + """ + Open a message and inform the GRIB file that it's been incremented. + + If ``grib_file`` is not supplied, the message is cloned from + ``GribMessage`` ``clone``. If neither is supplied, the ``GribMessage`` + is cloned from ``sample``. If ``index`` is supplied as a GribIndex, the + message is taken from the index. + """ + #: Unique GRIB ID, for GRIB API interface + self.gid = None + #: File containing message + self.grib_file = None + #: GribIndex referencing message + self.grib_index = None + if grib_file is not None: + self.gid = eccodes.codes_grib_new_from_file(grib_file.file_handle) + if self.gid is None: + raise IOError("Grib file %s is exhausted" % grib_file.name) + self.grib_file = grib_file + self.grib_file.message += 1 + self.grib_file.open_messages.append(self) + elif clone is not None: + self.gid = eccodes.codes_clone(clone.gid) + elif sample is not None: + self.gid = eccodes.codes_grib_new_from_samples(sample) + elif gribindex is not None: + self.gid = eccodes.codes_new_from_index(gribindex.iid) + if not self.gid: + raise IndexNotSelectedError("All keys must have selected " + "values before receiving message " + "from index.") + self.grib_index = gribindex + gribindex.open_messages.append(self) + else: + raise RuntimeError("Either grib_file, clone, sample or gribindex " + "must be provided.") + + def size(self): + """Return size of message in bytes.""" + return eccodes.codes_get_message_size(self.gid) + + def dump(self): + """Dump message's binary content.""" + return eccodes.codes_get_message(self.gid) + + def get(self, key, ktype=None): + """Get value of a given key as its native or specified type.""" + if self.missing(key): + raise KeyError("Key is missing from message.") + if eccodes.codes_get_size(self.gid, key) > 1: + ret = eccodes.codes_get_array(self.gid, key, ktype) + else: + ret = eccodes.codes_get(self.gid, key, ktype) + return ret + + def missing(self, key): + """Report if key is missing.""" + return bool(eccodes.codes_is_missing(self.gid, key)) + + def set_missing(self, key): + """Set a key to missing.""" + eccodes.codes_set_missing(self.gid, key) + + def write(self, outfile=None): + """Write message to file.""" + if not outfile: + # This is a hack because the API does not accept inheritance + outfile = self.grib_file.file_handle + eccodes.codes_write(self.gid, outfile) diff --git a/python/setup.py.in b/python/setup.py.in index 9de76e2a9..09eed7dde 100644 --- a/python/setup.py.in +++ b/python/setup.py.in @@ -68,4 +68,4 @@ setup(name='eccodes', url='https://software.ecmwf.int/wiki/display/ECC/ecCodes+Home', download_url='https://software.ecmwf.int/wiki/display/ECC/Releases', ext_modules=[Extension('gribapi._gribapi_swig', **attdict)], - packages=['eccodes', 'gribapi']) + packages=['eccodes', 'eccodes.high_level', 'gribapi'])