From db43d56513939e1e3ac4db97e3a6e85491aaaa66 Mon Sep 17 00:00:00 2001
From: Patrick Peglar <patrick.peglar@metoffice.gov.uk>
Date: Wed, 1 May 2024 12:18:16 +0100
Subject: [PATCH] Tidy code, do routine introspection just once.

---
 src/iris_grib/__init__.py | 63 +++++++++++++++++++++------------------
 src/iris_grib/message.py  |  4 +--
 2 files changed, 36 insertions(+), 31 deletions(-)

diff --git a/src/iris_grib/__init__.py b/src/iris_grib/__init__.py
index e821e0b9..c9265e06 100644
--- a/src/iris_grib/__init__.py
+++ b/src/iris_grib/__init__.py
@@ -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.
@@ -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)
@@ -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
diff --git a/src/iris_grib/message.py b/src/iris_grib/message.py
index 64cd2a90..6f5ad948 100644
--- a/src/iris_grib/message.py
+++ b/src/iris_grib/message.py
@@ -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