mirror of https://github.com/ecmwf/eccodes.git
Implement CodesMessage abstract class
This commit is contained in:
parent
613f5147ef
commit
a79e5cf2e5
|
@ -11,33 +11,10 @@ from .. import eccodes
|
|||
class CodesFile(file):
|
||||
|
||||
"""
|
||||
An abstract class to specify and/or implement behavior that files read by
|
||||
ecCodes should implement.
|
||||
An abstract class to specify and/or implement common behavior that files
|
||||
read by ecCodes should implement.
|
||||
"""
|
||||
|
||||
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 number of messages in file."""
|
||||
return eccodes.codes_count_in_file(self.file_handle)
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def next(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def __init__(self, filename, mode="r"):
|
||||
"""Open file and receive codes file handle."""
|
||||
#: File handle for working with actual file on disc
|
||||
|
@ -48,3 +25,26 @@ class CodesFile(file):
|
|||
self.message = 0
|
||||
#: Open messages
|
||||
self.open_messages = []
|
||||
|
||||
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 __len__(self):
|
||||
"""Return total number of messages in file."""
|
||||
return eccodes.codes_count_in_file(self.file_handle)
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def close(self):
|
||||
"""Possibility to manually close file."""
|
||||
self.__exit__(None, None, None)
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def next(self):
|
||||
raise NotImplementedError
|
||||
|
|
|
@ -0,0 +1,152 @@
|
|||
"""
|
||||
``CodesMessage`` class that implements a message readable by ecCodes 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, 2016
|
||||
"""
|
||||
|
||||
from .. import eccodes
|
||||
|
||||
|
||||
class CodesMessage(object):
|
||||
|
||||
"""
|
||||
An abstract class to specify and/or implement common behavior that
|
||||
messages read by ecCodes should implement.
|
||||
"""
|
||||
|
||||
def __init__(self, codes_file=None, clone=None, sample=None,
|
||||
other_args_found=False):
|
||||
"""
|
||||
Open a message and inform the host file that it's been incremented.
|
||||
|
||||
If ``codes_file`` is not supplied, the message is cloned from
|
||||
``CodesMessage`` ``clone``. If neither is supplied,
|
||||
the ``CodesMessage`` is cloned from ``sample``.
|
||||
|
||||
:param codes_file: A file readable for ecCodes
|
||||
:param clone: A valid ``CodesMessage``
|
||||
:param sample: A valid sample path to create ``CodesMessage`` from
|
||||
"""
|
||||
if not other_args_found and codes_file is None and clone is None and \
|
||||
sample is None:
|
||||
raise RuntimeError("CodesMessage initialization parameters not "
|
||||
"present.")
|
||||
#: Unique ID, for ecCodes interface
|
||||
self.codes_id = None
|
||||
#: File containing message
|
||||
self.codes_file = None
|
||||
if codes_file is not None:
|
||||
self.codes_id = self.new_from_file(codes_file.file_handle)
|
||||
if self.codes_id is None:
|
||||
raise IOError("CodesFile %s is exhausted" % codes_file.name)
|
||||
self.codes_file = codes_file
|
||||
self.codes_file.message += 1
|
||||
self.codes_file.open_messages.append(self)
|
||||
elif clone is not None:
|
||||
self.codes_id = eccodes.codes_clone(clone.codes_id)
|
||||
elif sample is not None:
|
||||
self.codes_id = self.new_from_sample(sample)
|
||||
|
||||
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.codes_file.file_handle
|
||||
eccodes.codes_write(self.codes_id, outfile)
|
||||
|
||||
@staticmethod
|
||||
def new_from_sample(samplename):
|
||||
"""For generic handling of message types."""
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def new_from_file(fileobj, headers_only=False):
|
||||
"""For generic handling of message types."""
|
||||
raise NotImplementedError
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
"""
|
||||
Set value associated with key.
|
||||
|
||||
Iterables and scalars are handled intelligently.
|
||||
"""
|
||||
# Passed value is iterable and not string
|
||||
if hasattr(value, "__iter__"):
|
||||
eccodes.codes_set_array(self.codes_id, key, value)
|
||||
else:
|
||||
eccodes.codes_set(self.codes_id, key, value)
|
||||
|
||||
def keys(self, namespace=None):
|
||||
"""Get available keys in message."""
|
||||
iterator = eccodes.codes_keys_iterator_new(self.codes_id,
|
||||
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 size(self):
|
||||
"""Return size of message in bytes."""
|
||||
return eccodes.codes_get_message_size(self.codes_id)
|
||||
|
||||
def dump(self):
|
||||
"""Dump message's binary content."""
|
||||
return eccodes.codes_get_message(self.codes_id)
|
||||
|
||||
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.codes_id, key) > 1:
|
||||
ret = eccodes.codes_get_array(self.codes_id, key, ktype)
|
||||
else:
|
||||
ret = eccodes.codes_get(self.codes_id, key, ktype)
|
||||
return ret
|
||||
|
||||
def missing(self, key):
|
||||
"""Report if key is missing."""
|
||||
return bool(eccodes.codes_is_missing(self.codes_id, key))
|
||||
|
||||
def set_missing(self, key):
|
||||
"""Set a key to missing."""
|
||||
eccodes.codes_set_missing(self.codes_id, key)
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Release message handle and inform host file of release."""
|
||||
eccodes.codes_release(self.codes_id)
|
||||
|
||||
def __enter__(self):
|
||||
return 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 __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()]
|
||||
|
|
@ -6,8 +6,7 @@ message when it is no longer needed, coordinating this with its host file.
|
|||
Author: Daniel Lee, DWD, 2014
|
||||
"""
|
||||
|
||||
import collections
|
||||
|
||||
from .codesmessage import CodesMessage
|
||||
from .. import eccodes
|
||||
|
||||
|
||||
|
@ -15,7 +14,8 @@ class IndexNotSelectedError(Exception):
|
|||
"""GRIB index was requested before selecting key/value pairs."""
|
||||
|
||||
|
||||
class GribMessage(object):
|
||||
class GribMessage(CodesMessage):
|
||||
|
||||
"""
|
||||
A GRIB message.
|
||||
|
||||
|
@ -59,140 +59,63 @@ class GribMessage(object):
|
|||
... # 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):
|
||||
# Arguments included explicitly to support introspection
|
||||
# TODO: Include headers_only option
|
||||
def __init__(self, codes_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.
|
||||
The message is taken from ``codes_file``, cloned from ``clone`` or
|
||||
``sample``, or taken from ``index``, in that order of precedence.
|
||||
"""
|
||||
#: Unique GRIB ID, for GRIB API interface
|
||||
self.gid = None
|
||||
#: File containing message
|
||||
self.grib_file = None
|
||||
grib_args_present = True
|
||||
if gribindex is None:
|
||||
grib_args_present = False
|
||||
super(self.__class__, self).__init__(codes_file, clone, sample,
|
||||
grib_args_present)
|
||||
#: 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:
|
||||
if gribindex is not None:
|
||||
self.codes_id = eccodes.codes_new_from_index(gribindex.iid)
|
||||
if not self.codes_id:
|
||||
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 __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Release GRIB message handle and inform file of release."""
|
||||
super(self.__class__, self).__exit__(exc_type, exc_val, exc_tb)
|
||||
if self.grib_index:
|
||||
self.grib_index.open_messages.remove(self)
|
||||
|
||||
def dump(self):
|
||||
"""Dump message's binary content."""
|
||||
return eccodes.codes_get_message(self.gid)
|
||||
# I'd like these to be simple function mappings but this worked better in
|
||||
# the interest of time
|
||||
@staticmethod
|
||||
def new_from_sample(samplename):
|
||||
return eccodes.codes_grib_new_from_samples(samplename)
|
||||
|
||||
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
|
||||
@staticmethod
|
||||
def new_from_file(fileobj, headers_only=False):
|
||||
return eccodes.codes_grib_new_from_file(fileobj, headers_only)
|
||||
|
||||
def missing(self, key):
|
||||
"""Report if key is missing."""
|
||||
return bool(eccodes.codes_is_missing(self.gid, key))
|
||||
@property
|
||||
def gid(self):
|
||||
"""Provided for backwards compatibility."""
|
||||
return self.codes_id
|
||||
|
||||
def set_missing(self, key):
|
||||
"""Set a key to missing."""
|
||||
eccodes.codes_set_missing(self.gid, key)
|
||||
@property
|
||||
def grib_file(self):
|
||||
"""Provided for backwards compatibility."""
|
||||
return self.codes_file
|
||||
|
||||
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)
|
||||
@gid.setter
|
||||
def gid(self, val):
|
||||
self.codes_id = val
|
||||
|
||||
@grib_file.setter
|
||||
def grib_file(self, val):
|
||||
self.codes_file = val
|
||||
|
|
Loading…
Reference in New Issue