SciTools · lbdreyer · Dec 4, 2019 · Sep 6, 2019 · Sep 13, 2019 · Sep 13, 2019
diff --git a/iris_grib/_load_convert.py b/iris_grib/_load_convert.py
@@ -26,13 +26,13 @@
 from iris.coords import AuxCoord, DimCoord, CellMethod
 from iris.exceptions import TranslationError
 from . import grib_phenom_translation as itranslation
+from .grib_phenom_translation import GRIBCode
 from iris.fileformats.rules import ConversionMetadata, Factory, Reference, \
     ReferenceTarget
 from iris.util import _is_circular
 
 from ._iris_mercator_support import confirm_extended_mercator_supported
 from ._grib1_load_rules import grib1_convert
-from .message import GribMessage
 
 
 # Restrict the names imported from this namespace.
@@ -1378,6 +1378,13 @@ def translate_phenomenon(metadata, discipline, parameterCategory,
             metadata['long_name'] = long_name
             metadata['units'] = Unit(1)
 
+    # Add a standard attribute recording the grib phenomenon identity.
+    metadata['attributes']['GRIB_PARAM'] = GRIBCode(
+        edition_or_string=2,
+        discipline=discipline,
+        category=parameterCategory,
+        number=parameterNumber)
+
     # Identify hybrid height and pressure reference fields.
     # Look for fields at surface level first.
     if (typeOfFirstFixedSurface == 1 and

diff --git a/iris_grib/_save_rules.py b/iris_grib/_save_rules.py
@@ -30,6 +30,7 @@
 from . import grib_phenom_translation as gptx
 from ._load_convert import (_STATISTIC_TYPE_NAMES, _TIME_RANGE_UNITS,
                             _SPATIAL_PROCESSING_TYPES)
+from .grib_phenom_translation import GRIBCode
 from iris.util import is_regular, regular_step
 
 
@@ -690,22 +691,45 @@ def grid_definition_section(cube, grib):
 ###############################################################################
 
 def set_discipline_and_parameter(cube, grib):
-    # NOTE: for now, can match by *either* standard_name or long_name.
-    # This allows workarounds for data with no identified standard_name.
-    grib2_info = gptx.cf_phenom_to_grib2_info(cube.standard_name,
-                                              cube.long_name)
-    if grib2_info is not None:
-        gribapi.grib_set(grib, "discipline", grib2_info.discipline)
-        gribapi.grib_set(grib, "parameterCategory", grib2_info.category)
-        gribapi.grib_set(grib, "parameterNumber", grib2_info.number)
-    else:
-        gribapi.grib_set(grib, "discipline", 255)
-        gribapi.grib_set(grib, "parameterCategory", 255)
-        gribapi.grib_set(grib, "parameterNumber", 255)
+    # Default values for parameter identity keys = effectively "MISSING".
+    discipline, category, number = 255, 255, 255
+    identity_found = False
+
+    # First, see if we can find and interpret a 'GRIB_PARAM' attribute.
+    attr = cube.attributes.get('GRIB_PARAM', None)
+    if attr:
+        try:
+            # Convert to standard tuple-derived form.
+            gc = GRIBCode(attr)
+            if gc.edition == 2:
+                discipline = gc.discipline
+                category = gc.category
+                number = gc.number
+                identity_found = True
+        except:
+            pass
+
+    if not identity_found:
+        # Else, translate a cube phenomenon, if possible.
+        # NOTE: for now, can match by *either* standard_name or long_name.
+        # This allows workarounds for data with no identified standard_name.
+        grib2_info = gptx.cf_phenom_to_grib2_info(cube.standard_name,
+                                                  cube.long_name)
+        if grib2_info is not None:
+            discipline = grib2_info.discipline
+            category = grib2_info.category
+            number = grib2_info.number
+            identity_found = True
+
+    if not identity_found:
         warnings.warn('Unable to determine Grib2 parameter code for cube.\n'
                       'discipline, parameterCategory and parameterNumber '
                       'have been set to "missing".')
 
+    gribapi.grib_set(grib, "discipline", discipline)
+    gribapi.grib_set(grib, "parameterCategory", category)
+    gribapi.grib_set(grib, "parameterNumber", number)
+
 
 def _non_missing_forecast_period(cube):
     # Calculate "model start time" to use as the reference time.

diff --git a/iris_grib/grib_phenom_translation.py b/iris_grib/grib_phenom_translation.py
@@ -16,8 +16,8 @@
 * cf --> grib2
 
 '''
-
-import collections
+from collections import namedtuple
+import re
 import warnings
 
 import cf_units
@@ -54,12 +54,12 @@ def __setitem__(self, key, value):
 
 # Define namedtuples for keys+values of the Grib1 lookup table.
 
-_Grib1ToCfKeyClass = collections.namedtuple(
+_Grib1ToCfKeyClass = namedtuple(
     'Grib1CfKey',
     ('table2_version', 'centre_number', 'param_number'))
 
 # NOTE: this form is currently used for both Grib1 *and* Grib2
-_GribToCfDataClass = collections.namedtuple(
+_GribToCfDataClass = namedtuple(
     'Grib1CfData',
     ('standard_name', 'long_name', 'units', 'set_height'))
 
@@ -146,7 +146,7 @@ def _make_grib1_cf_entry(table2_version, centre_number, param_number,
 
 # Define a namedtuple for the keys of the Grib2 lookup table.
 
-_Grib2ToCfKeyClass = collections.namedtuple(
+_Grib2ToCfKeyClass = namedtuple(
     'Grib2CfKey',
     ('param_discipline', 'param_category', 'param_number'))
 
@@ -205,11 +205,11 @@ def _make_grib2_cf_entry(param_discipline, param_category, param_number,
 
 # Define namedtuples for key+values of the cf-to-grib2 lookup table.
 
-_CfToGrib2KeyClass = collections.namedtuple(
+_CfToGrib2KeyClass = namedtuple(
     'CfGrib2Key',
     ('standard_name', 'long_name'))
 
-_CfToGrib2DataClass = collections.namedtuple(
+_CfToGrib2DataClass = namedtuple(
     'CfGrib2Data',
     ('discipline', 'category', 'number', 'units'))
 
@@ -316,3 +316,70 @@ def cf_phenom_to_grib2_info(standard_name, long_name=None):
     if standard_name is not None:
         long_name = None
     return _CF_GRIB2_TABLE[(standard_name, long_name)]
+
+
+class GRIBCode(namedtuple('GRIBCode',
+                          'edition discipline category number')):
+    """
+    An object representing a specific Grib phenomenon identity.
+
+    Basically a namedtuple of (edition, discipline, category, number).
+
+    Also provides a string representation, and supports creation from: another
+    similar object; a tuple of numbers; or any string with 4 separate decimal
+    numbers in it.
+
+    """
+    __slots__ = ()
+
+    def __new__(cls, edition_or_string,
+                discipline=None, category=None, number=None):
+        args = (edition_or_string, discipline, category, number)
+        nargs = sum(arg is not None for arg in args)
+        if nargs == 1:
+            # Single argument: convert to a string and extract 4 integers.
+            # NOTE: this also allows input from a GRIBCode, or a plain tuple.
+            edition_or_string = str(edition_or_string)
+            edition, discipline, category, number = \
+                cls._fournums_from_gribcode_string(edition_or_string)
+        elif nargs == 4:
+            edition = edition_or_string
+            edition, discipline, category, number = [
+                int(arg)
+                for arg in (edition, discipline, category, number)]
+        else:
+            msg = ('Cannot create GRIBCode from {} arguments, '
+                   '"GRIBCode{!r}" : '
+                   'expected either 1 or 4 non-None arguments.')
+            raise ValueError(msg.format(nargs, args))
+
+        return super(GRIBCode, cls).__new__(
+            cls, edition, discipline, category, number)
+
+    RE_PARSE_FOURNUMS = re.compile(4 * r'[^\d]*(\d*)')
+
+    @classmethod
+    def _fournums_from_gribcode_string(cls, edcn_string):
+        parsed_ok = False
+        nums_match = cls.RE_PARSE_FOURNUMS.match(edcn_string).groups()
+        if nums_match is not None:
+            try:
+                nums = [int(grp) for grp in nums_match]
+                parsed_ok = True
+            except ValueError:
+                pass
+
+        if not parsed_ok:
+            msg = ('Invalid argument for GRIBCode creation, '
+                   '"GRIBCode({!r})" : '
+                   'requires 4 numbers, separated by non-numerals.')
+            raise ValueError(msg.format(edcn_string))
+
+        return nums
+
+    PRINT_FORMAT = 'GRIB{:1d}:d{:03d}c{:03d}n{:03d}'
+
+    def __str__(self):
+        result = self.PRINT_FORMAT.format(
+            self.edition, self.discipline, self.category, self.number)
+        return result
diff --git a/iris_grib/tests/results/unit/load_cubes/load_cubes/reduced_raw.cml b/iris_grib/tests/results/unit/load_cubes/load_cubes/reduced_raw.cml
@@ -2,6 +2,7 @@
 <cubes xmlns="urn:x-iris:cubeml-0.2">
   <cube dtype="float64" standard_name="geopotential" units="m2 s-2">
     <attributes>
+      <attribute name="GRIB_PARAM" value="GRIB2:d000c003n004"/>
       <attribute name="centre" value="European Centre for Medium Range Weather Forecasts"/>
     </attributes>
     <coords>

diff --git a/iris_grib/tests/unit/grib_phenom_translation/test_grib_phenom_translation.py b/iris_grib/tests/unit/grib_phenom_translation/test_grib_phenom_translation.py
@@ -18,6 +18,7 @@
 import cf_units
 
 import iris_grib.grib_phenom_translation as gptx
+from iris_grib.grib_phenom_translation import GRIBCode
 
 
 class TestGribLookupTableType(tests.IrisTest):
@@ -151,5 +152,113 @@ def check_cf_grib2(standard_name, long_name,
                        expect_none=True)
 
 
+class TestGRIBcode(tests.IrisTest):
+    # GRIBCode is basically a namedtuple, so not all behaviour needs testing.
+    # However, creation is a bit special so exercise all those cases.
+    def test_create_from_keys(self):
+        gribcode = GRIBCode(
+            edition_or_string=5,
+            discipline=7,
+            category=4,
+            number=199)
+        self.assertEqual(gribcode.edition, 5)
+        self.assertEqual(gribcode.discipline, 7)
+        self.assertEqual(gribcode.category, 4)
+        self.assertEqual(gribcode.number, 199)
+
+    def test_create_from_args(self):
+        gribcode = GRIBCode(7, 3, 12, 99)
+        self.assertEqual(gribcode.edition, 7)
+        self.assertEqual(gribcode.discipline, 3)
+        self.assertEqual(gribcode.category, 12)
+        self.assertEqual(gribcode.number, 99)
+
+    def test_create_is_copy(self):
+        gribcode1 = GRIBCode(7, 3, 12, 99)
+        gribcode2 = GRIBCode(7, 3, 12, 99)
+        self.assertEqual(gribcode1, gribcode2)
+        self.assertIsNot(gribcode1, gribcode2)
+
+    def test_create_from_gribcode(self):
+        gribcode1 = GRIBCode((4, 3, 2, 1))
+        gribcode2 = GRIBCode(gribcode1)
+        self.assertEqual(gribcode1, gribcode2)
+        # NOTE: *not* passthrough : it creates a copy
+        # (though maybe not too significant, as it is immutable anyway?)
+        self.assertIsNot(gribcode1, gribcode2)
+
+    def test_create_from_string(self):
+        gribcode = GRIBCode('xxx12xs-34 -5,678qqqq')
+        # NOTE: args 2 and 3 are *not* negative.
+        self.assertEqual(gribcode, GRIBCode(12, 34, 5, 678))
+
+    def test_create_from_own_string(self):
+        # Check that GRIBCode string reprs are valid as create arguments.
+        gribcode = GRIBCode(
+            edition_or_string=2,
+            discipline=17,
+            category=94,
+            number=231)
+        grib_param_string = str(gribcode)
+        newcode = GRIBCode(grib_param_string)
+        self.assertEqual(newcode, gribcode)
+
+    def test_create_from_tuple(self):
+        gribcode = GRIBCode((4, 3, 2, 1))
+        self.assertEqual(gribcode, GRIBCode(4, 3, 2, 1))
+
+    def test_create_bad_nargs(self):
+        # Between 1 and 4 args is not invalid call syntax, but it should fail.
+        with self.assertRaisesRegex(
+                ValueError,
+                'Cannot create GRIBCode from 2 arguments'):
+            GRIBCode(1, 2)
+
+    def test_create_bad_single_arg_None(self):
+        with self.assertRaisesRegex(
+                ValueError,
+                'Cannot create GRIBCode from 0 arguments'):
+            GRIBCode(None)
+
+    def test_create_bad_single_arg_empty_string(self):
+        with self.assertRaisesRegex(
+                ValueError,
+                'Invalid argument for GRIBCode creation'):
+            GRIBCode('')
+
+    def test_create_bad_single_arg_nonums(self):
+        with self.assertRaisesRegex(
+                ValueError,
+                'Invalid argument for GRIBCode creation'):
+            GRIBCode('saas- dsa- ')
+
+    def test_create_bad_single_arg_less_than_4_nums(self):
+        with self.assertRaisesRegex(
+                ValueError,
+                'Invalid argument for GRIBCode creation'):
+            GRIBCode('1,2,3')
+
+    def test_create_bad_single_arg_number(self):
+        with self.assertRaisesRegex(
+                ValueError,
+                'Invalid argument for GRIBCode creation'):
+            GRIBCode(4)
+
+    def test_create_bad_single_arg_single_numeric(self):
+        with self.assertRaisesRegex(
+                ValueError,
+                'Invalid argument for GRIBCode creation'):
+            GRIBCode('44')
+
+    def test_create_string_more_than_4_nums(self):
+        # Note: does not error, just discards the extra.
+        gribcode = GRIBCode('1,2,3,4,5,6,7,8')
+        self.assertEqual(gribcode, GRIBCode(1, 2, 3, 4))
+
+    def test__str__(self):
+        result = str(GRIBCode(2, 17, 3, 123))
+        self.assertEqual(result, 'GRIB2:d017c003n123')
+
+
 if __name__ == '__main__':
     tests.main()
diff --git a/iris_grib/tests/unit/load_convert/test_translate_phenomenon.py b/iris_grib/tests/unit/load_convert/test_translate_phenomenon.py
@@ -18,7 +18,7 @@
 from iris.coords import DimCoord
 
 from iris_grib._load_convert import Probability, translate_phenomenon
-from iris_grib.grib_phenom_translation import _GribToCfDataClass
+from iris_grib.grib_phenom_translation import _GribToCfDataClass, GRIBCode
 
 
 class Test_probability(tests.IrisGribTest):
@@ -30,11 +30,11 @@ def setUp(self):
             return_value=_GribToCfDataClass('air_temperature', '', 'K', None))
         # Construct dummy call arguments
         self.probability = Probability('<prob_type>', 22.0)
-        self.metadata = {'aux_coords_and_dims': []}
+        self.metadata = {'aux_coords_and_dims': [], 'attributes': {}}
 
     def test_basic(self):
-        result = translate_phenomenon(self.metadata, None, None, None, None,
-                                      None, None, probability=self.probability)
+        translate_phenomenon(self.metadata, 7, 8, 9, None,
+                             None, None, probability=self.probability)
         # Check metadata.
         thresh_coord = DimCoord([22.0],
                                 standard_name='air_temperature',
@@ -43,14 +43,23 @@ def test_basic(self):
             'standard_name': None,
             'long_name': 'probability_of_air_temperature_<prob_type>',
             'units': Unit(1),
-            'aux_coords_and_dims': [(thresh_coord, None)]})
+            'aux_coords_and_dims': [(thresh_coord, None)],
+            'attributes': {'GRIB_PARAM': GRIBCode(2, 7, 8, 9)}})
 
     def test_no_phenomenon(self):
-        original_metadata = deepcopy(self.metadata)
         self.phenom_lookup_patch.return_value = None
-        result = translate_phenomenon(self.metadata, None, None, None, None,
-                                      None, None, probability=self.probability)
-        self.assertEqual(self.metadata, original_metadata)
+        expected_metadata = self.metadata.copy()
+        translate_phenomenon(self.metadata,
+                             discipline=7,
+                             parameterCategory=77,
+                             parameterNumber=777,
+                             typeOfFirstFixedSurface=None,
+                             scaledValueOfFirstFixedSurface=None,
+                             typeOfSecondFixedSurface=None,
+                             probability=self.probability)
+        expected_metadata['attributes']['GRIB_PARAM'] = \
+            GRIBCode(2, 7, 77, 777)
+        self.assertEqual(self.metadata, expected_metadata)
 
 
 if __name__ == '__main__':