Pull request #3: Feature/set multiple

Merge in ECCODES/eccodes-python from feature/set_multiple to develop

* commit 'b523e5abd4dd0b810bfe32079d05e55b9906e12f':
  Black formatting
  Improve value error message
  Improve value error message
  Fix black formatting using v22.1.0
  Formatting
  Modify set function in message to accept dictionary of key values and add check on values

Author: jinmannwong

eccodes/highlevel/message.py

@@ -67,16 +67,46 @@ def get(self, name, default=None, ktype=None):
         except KeyError:
             return default
 
-    def set(self, name, value):
-        """Set the value of the given key
+    def set(self, *args):
+        """If two arguments are given, assumes this takes form of a single key
+        value pair and sets the value of the given key. If a dictionary is passed
+        then sets the values of all keys in the dictionary. Note, ordering
+        if the keys is important. Finally, checks if value
+        has been set correctly
 
         Raises
         ------
+        TypeError
+            If arguments do not take one of the two expected forms
         KeyError
             If the key does not exist
+        ValueError
+            If the set value of one of the keys is not the expected value
         """
-        with raise_keyerror(name):
-            return eccodes.codes_set(self._handle, name, value)
+        if isinstance(args[0], str) and len(args) == 2:
+            key_values = {args[0]: args[1]}
+        elif isinstance(args[0], dict):
+            key_values = args[0]
+        else:
+            raise TypeError(
+                f"Unsupported argument type. Expects two arguments consisting \
+            of key and value pair, or a dictionary of key-value pairs"
+            )
+
+        for name, value in key_values.items():
+            with raise_keyerror(name):
+                eccodes.codes_set(self._handle, name, value)
+
+        # Check values just set
+        for name, value in key_values.items():
+            saved_value = self.get(name)
+            cast_value = value
+            if not isinstance(value, type(saved_value)):
+                cast_value = type(saved_value)(value)
+            if saved_value != cast_value:
+                raise ValueError(
+                    f"Unexpected retrieved value {saved_value} for key {name}. Expected {cast_value}"
+                )
 
     def get_array(self, name):
         """Get the value of the given key as an array

tests/test_highlevel.py

@@ -65,6 +65,23 @@ def test_message_set():
         assert message.is_missing(missing_key)
 
 
+def test_message_set_multiple():
+    with eccodes.FileReader(TEST_GRIB_DATA) as reader:
+        message = next(reader)
+        message.set(
+            {
+                "centre": "ecmf",
+                "numberOfValues": 10,
+                "shortName": "z",
+            }
+        )
+        with pytest.raises(TypeError):
+            message.set("centre", "ecmwf", 2)
+        with pytest.raises(ValueError):
+            message.set("stepRange", "0-12")
+        message.set({"stepType": "max", "stepRange": "0-12"})
+
+
 def test_message_iter():
     with eccodes.FileReader(TEST_GRIB_DATA2) as reader:
         message = next(reader)