Skip to content

Commit

Permalink
Tidy code, do routine introspection just once.
Browse files Browse the repository at this point in the history
  • Loading branch information
pp-mo authored and bjlittle committed May 2, 2024
1 parent a0dd005 commit db43d56
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 31 deletions.
63 changes: 34 additions & 29 deletions src/iris_grib/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,39 @@ def __setstate__(self, state):
setattr(self, key, value)


# Utility routines for the use of dask 'meta' in wrapping proxies
def _aslazydata_has_meta():
"""
Work out whether 'iris._lazy_data.as_lazy_data' takes a "meta" kwarg.
Up to Iris 3.8.0, "as_lazy_data" did not have a 'meta' keyword, but
since https://github.com/SciTools/iris/pull/5801, it now *requires* one,
if the wrapped object is anything other than a numpy or dask array.
"""
from inspect import signature
from iris._lazy_data import as_lazy_data

sig = signature(as_lazy_data)
return "meta" in sig.parameters


# Work this out just once.
_ASLAZYDATA_NEEDS_META = _aslazydata_has_meta()


def _make_dask_meta(shape, dtype, is_masked=True):
"""
Construct a dask 'meta' object for use in 'dask.array.from_array'.
A "meta" array is made from the dtype and shape of the array-like to be
wrapped, plus whether it will return masked or unmasked data.
"""
meta_shape = tuple([0 for _ in shape])
array_class = np.ma if is_masked else np
meta = array_class.zeros(meta_shape, dtype=dtype)
return meta


class GribWrapper:
"""
Contains a pygrib object plus some extra keys of our own.
Expand Down Expand Up @@ -198,7 +231,7 @@ def __init__(self, grib_message, grib_fh=None, has_bitmap=True):
dtype = np.dtype(float) # Use default dtype for python float
proxy = GribDataProxy(shape, dtype, grib_fh.name, offset)
as_lazy_kwargs = {}
if _aslazydata_has_meta():
if _ASLAZYDATA_NEEDS_META:
meta = _make_dask_meta(shape, dtype, is_masked=has_bitmap)
as_lazy_kwargs["meta"] = meta
self._data = as_lazy_data(proxy, **as_lazy_kwargs)
Expand Down Expand Up @@ -896,31 +929,3 @@ def save_messages(messages, target, append=False):
# (this bit is common to the pp and grib savers...)
if isinstance(target, str):
grib_file.close()


# Odd utility routines
def _aslazydata_has_meta():
"""
Work out whether 'iris._lazy_data.as_lazy_data' takes a "meta" kwarg.
Up to Iris 3.8.0, "as_lazy_data" did not have a 'meta' keyword, but
since https://github.com/SciTools/iris/pull/5801, it now *requires* one.
"""
from inspect import signature
from iris._lazy_data import as_lazy_data

sig = signature(as_lazy_data)
return "meta" in sig.parameters


def _make_dask_meta(shape, dtype, is_masked=True):
"""
Construct a dask 'meta' suitable for 'dask.array.from_array'.
Result is made from the dtype and shape of an array-like to be mapped,
plus whether it will return masked or unmasked arrays.
"""
meta_shape = tuple([0 for _ in shape])
array_class = np.ma if is_masked else np
meta = array_class.zeros(meta_shape, dtype=dtype)
return meta
4 changes: 2 additions & 2 deletions src/iris_grib/message.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,9 +171,9 @@ def data(self):
proxy = _DataProxy(shape, dtype, self._recreate_raw)

as_lazy_kwargs = {}
from . import _aslazydata_has_meta, _make_dask_meta
from . import _ASLAZYDATA_NEEDS_META, _make_dask_meta

if _aslazydata_has_meta():
if _ASLAZYDATA_NEEDS_META:
has_bitmap = 6 in sections
meta = _make_dask_meta(shape, dtype, is_masked=has_bitmap)
as_lazy_kwargs["meta"] = meta
Expand Down

0 comments on commit db43d56

Please sign in to comment.