Implement CodesMessage abstract class

This commit is contained in:
Daniel Lee 2016-12-16 14:31:20 +01:00
parent 613f5147ef
commit a79e5cf2e5
3 changed files with 222 additions and 147 deletions

View File

@ -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

View File

@ -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()]

View File

@ -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