Repository: ecmwf/earthkit-data (python)

Issue:
Allow None for  bits_per_value option in FieldList save() method
### Is your feature request related to a problem? Please describe.

Currently `save()` has the following options:

```python
def save(filename, append=False, check_nans=True, bits_per_value=16)
```
where `bits_per_value` must be an int (cannot be None).

This poses a problem when the Fieldlist was created by modifying another one, and we want keep the original fields' `bits_per_value` settings.

The simplest solution would be to allow `bits_per_value=None`, which would indicate to keep the original settings.

As for the default value, we should use:

```python
bits_per_value=None
```

### Describe the solution you'd like

_No response_

### Describe alternatives you've considered

_No response_

### Additional context

_No response_

### Organisation

_No response_

Pull request:
Allow None for bits_per_values when saving GRIB data
Fixes #326 

Patch:
diff --git a/earthkit/data/core/fieldlist.py b/earthkit/data/core/fieldlist.py
index 056d6b6..552be34 100644
--- a/earthkit/data/core/fieldlist.py
+++ b/earthkit/data/core/fieldlist.py
@@ -1177,6 +1177,13 @@ class FieldList(Index):
             the target file be overwritten if already exists. Default is False
         **kwargs: dict, optional
             Other keyword arguments passed to :obj:`write`.
+
+        See Also
+        --------
+        :obj:`write`
+        :meth:`GribFieldList.save() <data.readers.grib.index.GribFieldList.save>`
+        :meth:`NumpyFieldList.save() <data.sources.numpy_list.NumpyFieldList.save>`
+
         """
         flag = "wb" if not append else "ab"
         with open(filename, flag) as f:
@@ -1191,6 +1198,11 @@ class FieldList(Index):
             The target file object.
         **kwargs: dict, optional
             Other keyword arguments passed to the underlying field implementation.
+
+        See Also
+        --------
+        read
+
         """
         for s in self:
             s.write(f, **kwargs)
diff --git a/earthkit/data/core/metadata.py b/earthkit/data/core/metadata.py
index 990b964..b5f8c9c 100644
--- a/earthkit/data/core/metadata.py
+++ b/earthkit/data/core/metadata.py
@@ -33,13 +33,23 @@ class Metadata(metaclass=ABCMeta):
     INDEX_KEYS = []
     CUSTOM_KEYS = [DATETIME]
 
+    extra = None
+
     def __iter__(self):
         """Return an iterator over the metadata keys."""
         return iter(self.keys())
 
-    @abstractmethod
     def __len__(self):
         r"""Return the number of metadata entries."""
+        if not self.extra:
+            return self._len()
+        else:
+            extra = sum([1 for x in self.extra if not self._contains(x)])
+            return extra + self._len()
+
+    @abstractmethod
+    def _len(self):
+        r"""Return the number of metadata entries without the extra items."""
         pass
 
     def __getitem__(self, key):
@@ -52,7 +62,6 @@ class Metadata(metaclass=ABCMeta):
         """
         return self.get(key, raise_on_missing=True)
 
-    @abstractmethod
     def __contains__(self, key):
         r"""Check if ``key`` is available.
 
@@ -60,9 +69,23 @@ class Metadata(metaclass=ABCMeta):
         -------
         bool
         """
-        pass
+        if not self.extra:
+            return self._contains(key)
+        else:
+            if key in self.extra:
+                return True
+            return self._contains(key)
 
     @abstractmethod
+    def _contains(self, key):
+        r"""Check if ``key`` is available in the non extra-keys.
+
+        Returns
+        -------
+        bool
+        """
+        pass
+
     def keys(self):
         r"""Return the metadata keys.
 
@@ -71,9 +94,27 @@ class Metadata(metaclass=ABCMeta):
         Iterable of str
 
         """
-        pass
+        if not self.extra:
+            return self._keys()
+        else:
+            extra = [x for x in self.extra if x not in self._keys()]
+            if len(extra) == 0:
+                return self._keys()
+            else:
+                r = list(self._keys()) + extra
+                return r
 
     @abstractmethod
+    def _keys(self):
+        r"""Return the metadata keys without the extra keys.
+
+        Returns
+        -------
+        Iterable of str
+
+        """
+        pass
+
     def items(self):
         r"""Return the metadata items.
 
@@ -81,6 +122,22 @@ class Metadata(metaclass=ABCMeta):
         -------
         Iterable of :obj:`(key,value)` pairs
 
+        """
+        if not self.extra:
+            return self._items()
+        else:
+            r = dict(self._items())
+            r.update(self.extra)
+            return r.items()
+
+    @abstractmethod
+    def _items(self):
+        r"""Return the metadata items without the extra keys.
+
+        Returns
+        -------
+        Iterable of :obj:`(key,value)` pairs
+
         """
         pass
 
@@ -115,6 +172,8 @@ class Metadata(metaclass=ABCMeta):
             a missing value.
 
         """
+        if self._is_extra_key(key):
+            return self._get_extra_key(key, default=default, astype=astype)
         if self._is_custom_key(key):
             return self._get_custom_key(
                 key, default=default, astype=astype, raise_on_missing=raise_on_missing
@@ -128,6 +187,19 @@ class Metadata(metaclass=ABCMeta):
     def _get(self, key, astype=None, default=None, raise_on_missing=False):
         pass
 
+    def _is_extra_key(self, key):
+        return self.extra is not None and key in self.extra
+
+    def _get_extra_key(self, key, default=None, astype=None, **kwargs):
+        v = self.extra.get(key, default)
+
+        if astype is not None and v is not None:
+            try:
+                return astype(v)
+            except Exception:
+                return None
+        return v
+
     def _is_custom_key(self, key):
         return key in self.CUSTOM_KEYS and key not in self
 
@@ -300,9 +372,15 @@ class RawMetadata(Metadata):
     def __len__(self):
         return len(self._d)
 
+    def _len(self):
+        return self._len()
+
     def __contains__(self, key):
         return key in self._d
 
+    def _contains(self, key):
+        pass
+
     def _get(self, key, astype=None, default=None, raise_on_missing=False):
         if not raise_on_missing:
             v = self._d.get(key, default)
@@ -319,9 +397,15 @@ class RawMetadata(Metadata):
     def keys(self):
         return self._d.keys()
 
+    def _keys(self):
+        pass
+
     def items(self):
         return self._d.items()
 
+    def _items(self):
+        pass
+
     def _base_datetime(self):
         return None
 
diff --git a/earthkit/data/readers/grib/codes.py b/earthkit/data/readers/grib/codes.py
index 62a3b5c..34d6e53 100644
--- a/earthkit/data/readers/grib/codes.py
+++ b/earthkit/data/readers/grib/codes.py
@@ -291,16 +291,25 @@ class GribField(Field):
     #         name = "paramId"
     #     return self.handle.get(name)
 
-    def write(self, f):
+    def write(self, f, bits_per_value=None):
         r"""Writes the message to a file object.
 
         Parameters
         ----------
         f: file object
             The target file object.
+        bits_per_value: int or None
+            Set the ``bitsPerValue`` GRIB key in the generated GRIB message. When
+            None the ``bitsPerValue`` stored in the metadata will be used.
         """
+        if bits_per_value is not None:
+            handle = self.handle.clone()
+            handle.set_long("bitsPerValue", bits_per_value)
+        else:
+            handle = self.handle
+
         # assert isinstance(f, io.IOBase)
-        self.handle.write_to(f)
+        handle.write_to(f)
 
     def message(self):
         r"""Returns a buffer containing the encoded message.
diff --git a/earthkit/data/readers/grib/index/__init__.py b/earthkit/data/readers/grib/index/__init__.py
index bdb9128..d5020e4 100644
--- a/earthkit/data/readers/grib/index/__init__.py
+++ b/earthkit/data/readers/grib/index/__init__.py
@@ -14,7 +14,7 @@ from abc import abstractmethod
 
 from earthkit.data.core.fieldlist import FieldList
 from earthkit.data.core.index import Index, MaskIndex, MultiIndex
-from earthkit.data.decorators import alias_argument
+from earthkit.data.decorators import alias_argument, detect_out_filename
 from earthkit.data.indexing.database import (
     FILEPARTS_KEY_NAMES,
     MORE_KEY_NAMES,
@@ -179,6 +179,33 @@ class GribFieldList(PandasMixIn, XarrayMixIn, FieldList):
     def _normalize_kwargs_names(self, **kwargs):
         return kwargs
 
+    @detect_out_filename
+    def save(self, filename, append=False, bits_per_value=None):
+        r"""Write all the fields into a file.
+
+        Parameters
+        ----------
+        filename: str
+            The target file path.
+        append: bool
+            When it is true append data to the target file. Otherwise
+            the target file be overwritten if already exists.
+        bits_per_value: int or None
+            Set the ``bitsPerValue`` GRIB key for each message in the generated
+            output. When None the ``bitsPerValue`` stored in the message metadata
+            will be used.
+
+        See Also
+        --------
+        write
+
+        """
+        super().save(
+            filename,
+            append=append,
+            bits_per_value=bits_per_value,
+        )
+
 
 class GribMaskFieldList(GribFieldList, MaskIndex):
     def __init__(self, *args, **kwargs):
diff --git a/earthkit/data/readers/grib/metadata.py b/earthkit/data/readers/grib/metadata.py
index db1fbe9..16cbbef 100644
--- a/earthkit/data/readers/grib/metadata.py
+++ b/earthkit/data/readers/grib/metadata.py
@@ -11,6 +11,7 @@ import datetime
 
 from earthkit.data.core.geography import Geography
 from earthkit.data.core.metadata import Metadata
+from earthkit.data.decorators import cached_method
 from earthkit.data.indexing.database import GRIB_KEYS_NAMES
 from earthkit.data.utils.bbox import BoundingBox
 from earthkit.data.utils.projections import Projection
@@ -203,13 +204,15 @@ class GribMetadata(Metadata):
 
     __handle_type = None
 
-    def __init__(self, handle):
+    def __init__(self, handle, extra=None):
         if not isinstance(handle, self._handle_type()):
             raise TypeError(
                 f"GribMetadata: expected handle type {self._handle_type()}, got {type(handle)}"
             )
         self._handle = handle
         self._geo = None
+        if extra is not None:
+            self.extra = extra
 
     @staticmethod
     def _handle_type():
@@ -222,16 +225,16 @@ class GribMetadata(Metadata):
             GribMetadata.__handle_type = GribCodesHandle
         return GribMetadata.__handle_type
 
-    def __len__(self):
-        return sum(map(lambda i: 1, self.keys()))
+    def _len(self):
+        return sum(map(lambda i: 1, self._keys()))
 
-    def __contains__(self, key):
+    def _contains(self, key):
         return self._handle.__contains__(key)
 
-    def keys(self):
+    def _keys(self):
         return self._handle.keys()
 
-    def items(self):
+    def _items(self):
         return self._handle.items()
 
     def _get(self, key, astype=None, default=None, raise_on_missing=False):
@@ -254,8 +257,21 @@ class GribMetadata(Metadata):
     def override(self, *args, **kwargs):
         d = dict(*args, **kwargs)
         handle = self._handle.clone(headers_only=True)
+
+        # whether headers_only=True works depends on the eccCodes version and the
+        # message properties. We check it by comparing the message lengths.
+        shrunk = handle.get_long("totalLength") < self._handle.get_long("totalLength")
+
         handle.set_multiple(d)
-        return GribMetadata(handle)
+
+        # some keys, needed later, are not copied into the clone when
+        # headers_only=True. We store them as extra keys.
+        if shrunk:
+            extra = {"bitsPerValue": self._handle.get("bitsPerValue", default=0)}
+        else:
+            extra = None
+
+        return GribMetadata(handle, extra=extra)
 
     def as_namespace(self, namespace=None):
         r"""Return all the keys/values from a namespace.
@@ -364,6 +380,9 @@ class RestrictedGribMetadata(GribMetadata):
     """Hide internal keys and namespaces in GRIB metadata"""
 
     EKD_NAMESPACE = "grib"
+
+    # ideally bitsPerValue should be here. However, it is treated as an
+    # extra key and cannot be an internal key.
     INTERNAL_KEYS = [
         "minimum",
         "maximum",
@@ -382,7 +401,6 @@ class RestrictedGribMetadata(GribMetadata):
         "numberOfMissing",
         "numberOfCodedValues",
         "bitmapPresent",
-        "bitsPerValue",
         "offsetValuesBy",
         "packingError",
         "referenceValue",
@@ -393,13 +411,17 @@ class RestrictedGribMetadata(GribMetadata):
     INTERNAL_NAMESPACES = ["statistics"]
 
     def __init__(self, md):
-        super().__init__(md._handle)
+        super().__init__(md._handle, extra=md.extra)
 
-    def __len__(self):
+    @cached_method
+    def _len(self):
         if self.INTERNAL_KEYS:
-            return len(self.keys())
+            # NOTE: argumentless super cannot be used in list comprehension
+            m = super()._contains
+            surplus = [1 for x in self.INTERNAL_KEYS if not m(x)]
+            return super()._len() + len(surplus)
         else:
-            return super().__len__()
+            return super()._len()
 
     def _is_internal(self, key):
         ns, _, name = key.partition(".")
@@ -412,31 +434,31 @@ class RestrictedGribMetadata(GribMetadata):
         else:
             return name in self.INTERNAL_KEYS
 
-    def __contains__(self, key):
+    def _contains(self, key):
         if self.INTERNAL_KEYS:
-            return not self._is_internal(key) and super().__contains__(key)
+            return not self._is_internal(key) and super()._contains(key)
         else:
-            return super().__contains__(key)
+            return super()._contains(key)
 
-    def keys(self):
+    def _keys(self):
         if self.INTERNAL_KEYS:
             r = []
-            for k in super().keys():
+            for k in super()._keys():
                 if k not in self.INTERNAL_KEYS:
                     r.append(k)
             return r
         else:
-            return super().keys()
+            return super()._keys()
 
-    def items(self):
+    def _items(self):
         if self.INTERNAL_KEYS:
             r = {}
-            for k, v in super().items():
+            for k, v in super()._items():
                 if k not in self.INTERNAL_KEYS:
                     r[k] = v
             return r.items()
         else:
-            return super().items()
+            return super()._items()
 
     def get(self, key, default=None, *, astype=None, raise_on_missing=False):
         ns, _, name = key.partition(".")
diff --git a/earthkit/data/sources/numpy_list.py b/earthkit/data/sources/numpy_list.py
index 7a95994..caf5c08 100644
--- a/earthkit/data/sources/numpy_list.py
+++ b/earthkit/data/sources/numpy_list.py
@@ -133,7 +133,7 @@ class NumpyFieldListCore(PandasMixIn, XarrayMixIn, FieldList):
         else:
             return type(self)(self.to_numpy(**kwargs), self._metadata)
 
-    def save(self, filename, append=False, check_nans=True, bits_per_value=16):
+    def save(self, filename, append=False, check_nans=True, bits_per_value=None):
         r"""Write all the fields into a file.
 
         Parameters
@@ -145,8 +145,9 @@ class NumpyFieldListCore(PandasMixIn, XarrayMixIn, FieldList):
             the target file be overwritten if already exists.
         check_nans: bool
             Replace nans in the values with GRIB missing values when generating the output.
-        bits_per_value: int
-            Set the ``bitsPerValue`` GRIB key in the generated output.
+        bits_per_value: int or None
+            Set the ``bitsPerValue`` GRIB key in the generated output. When None the
+            ``bitsPerValue`` stored in the metadata will be used.
         """
         super().save(
             filename,
diff --git a/earthkit/data/writers/grib.py b/earthkit/data/writers/grib.py
index 44bf13a..2362ab5 100644
--- a/earthkit/data/writers/grib.py
+++ b/earthkit/data/writers/grib.py
@@ -13,7 +13,7 @@ from . import Writer
 class GribWriter(Writer):
     DATA_FORMAT = "grib"
 
-    def write(self, f, values, metadata, check_nans=True, bits_per_value=16):
+    def write(self, f, values, metadata, check_nans=True, bits_per_value=None):
         r"""Write a GRIB field to a file object.
 
         Parameters
@@ -26,12 +26,17 @@ class GribWriter(Writer):
             Metadata of the GRIB field.
         check_nans: bool
             Replace nans in ``values`` with GRIB missing values when writing to ``f``.
-        bits_per_value: int
-            Set the ``bitsPerValue`` GRIB key in the generated GRIB message.
+        bits_per_value: int or None
+            Set the ``bitsPerValue`` GRIB key in the generated GRIB message. When
+            None the ``bitsPerValue`` stored in the metadata will be used.
         """
         handle = metadata._handle.clone()
 
-        handle.set_long("bitsPerValue", bits_per_value)
+        if bits_per_value is None:
+            bits_per_value = metadata.get("bitsPerValue", 0)
+
+        if bits_per_value != 0:
+            handle.set_long("bitsPerValue", bits_per_value)
 
         if check_nans:
             import numpy as np


Test patch:
diff --git a/tests/core/test_metadata.py b/tests/core/test_metadata.py
index 7e752d9..4c296fc 100644
--- a/tests/core/test_metadata.py
+++ b/tests/core/test_metadata.py
@@ -257,6 +257,40 @@ def test_grib_metadata_override_invalid():
     assert "EncodingError" in e.typename
 
 
+def test_grib_metadata_override_extra():
+    ds = from_source("file", earthkit_examples_file("test.grib"))
+    md = ds[0].metadata()
+    md_num = len(md)
+
+    assert md["perturbationNumber"] == 0
+    assert md["shortName"] == "2t"
+
+    extra = {"my_custom_key": "2", "shortName": "N", "perturbationNumber": 2}
+    md = GribMetadata(md._handle, extra=extra)
+
+    assert md["my_custom_key"] == "2"
+    assert md["perturbationNumber"] == 2
+    assert md["shortName"] == "N"
+    assert md["typeOfLevel"] == "surface"
+
+    keys = md.keys()
+    assert len(keys) == md_num + 1
+
+    for k in md.keys():
+        assert isinstance(k, str)
+        assert k != ""
+        break
+
+    items = md.items()
+    assert len(items) == md_num + 1
+
+    for k, v in md.items():
+        assert isinstance(k, str)
+        assert k != ""
+        assert v is not None
+        break
+
+
 if __name__ == "__main__":
     from earthkit.data.testing import main
 
diff --git a/tests/grib/test_grib_output.py b/tests/grib/test_grib_output.py
index a7ffa27..befe4e9 100644
--- a/tests/grib/test_grib_output.py
+++ b/tests/grib/test_grib_output.py
@@ -34,6 +34,19 @@ def test_grib_save_when_loaded_from_file():
         assert len(fs) == len(fs_saved)
 
 
+@pytest.mark.parametrize(
+    "_kwargs,expected_value",
+    [({}, 16), ({"bits_per_value": 12}, 12), ({"bits_per_value": None}, 16)],
+)
+def test_grib_save_bits_per_value(_kwargs, expected_value):
+    ds = from_source("file", earthkit_examples_file("test.grib"))
+
+    with temp_file() as tmp:
+        ds.save(tmp, **_kwargs)
+        ds1 = from_source("file", tmp)
+        assert ds1.metadata("bitsPerValue") == [expected_value] * len(ds)
+
+
 @pytest.mark.skipif(
     sys.version_info < (3, 10),
     reason="ignore_cleanup_errors requires Python 3.10 or later",
diff --git a/tests/numpy_fs/test_numpy_fs_metadata.py b/tests/numpy_fs/test_numpy_fs_metadata.py
index 52672a1..d86d90f 100644
--- a/tests/numpy_fs/test_numpy_fs_metadata.py
+++ b/tests/numpy_fs/test_numpy_fs_metadata.py
@@ -44,6 +44,9 @@ def test_numpy_fs_values_metadata():
         with pytest.raises(KeyError):
             ds[0].metadata(k)
 
+    # bits per value must be kept from the original GRIB data
+    assert ds[0].metadata("bitsPerValue") == 16
+
 
 def test_numpy_fs_values_metadata_internal():
     ds, _ = load_numpy_fs(1)
diff --git a/tests/numpy_fs/test_numpy_fs_write.py b/tests/numpy_fs/test_numpy_fs_write.py
index 5b73a1d..a6cbed8 100644
--- a/tests/numpy_fs/test_numpy_fs_write.py
+++ b/tests/numpy_fs/test_numpy_fs_write.py
@@ -150,11 +150,15 @@ def test_numpy_fs_grib_write_generating_proc_id():
 
 
 @pytest.mark.parametrize(
-    "_kwargs,expected_value", [({}, 16), ({"bits_per_value": 12}, 12)]
+    "_kwargs,expected_value",
+    [({}, None), ({"bits_per_value": 12}, 12), ({"bits_per_value": None}, None)],
 )
 def test_numpy_fs_grib_write_bits_per_value(_kwargs, expected_value):
     ds, _ = load_numpy_fs(1)
 
+    if expected_value is None:
+        expected_value = ds[0].metadata("bitsPerValue")
+
     with temp_file() as tmp:
         ds.save(tmp, **_kwargs)
         ds1 = from_source("file", tmp)