From cfa24b8c90a8cab54559d450c4a717bacacf730c Mon Sep 17 00:00:00 2001
From: Eddy Comyn-Platt <53045993+EddyCMWF@users.noreply.github.com>
Date: Thu, 2 Mar 2023 11:28:04 +0000
Subject: [PATCH] Feature/to netcdf var encoding (#2)

* var_encoding options for to_netcdf convertor
---
 cfgrib/__main__.py             | 25 ++++++++++++++++++++-----
 cfgrib/xarray_plugin.py        |  1 -
 tests/test_60_main_commands.py | 20 ++++++++++++++++++++
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/cfgrib/__main__.py b/cfgrib/__main__.py
index f8a8fe6a..6f34bfe5 100644
--- a/cfgrib/__main__.py
+++ b/cfgrib/__main__.py
@@ -78,13 +78,22 @@ def selfcheck() -> None:
     "-n",
     default=None,
     help=(
-        "kwargs used xarray.to_netcdf when creating the netCDF file."
-        "Can either be a JSON format string or "
-        "the path to JSON file."
+        "kwargs used xarray.to_netcdf when creating the netCDF file. "
+        "Can either be a JSON format string or the path to JSON file. "
     ),
 )
-def to_netcdf(inpaths, outpath, cdm, engine, backend_kwargs_json, netcdf_kwargs_json):
-    # type: (T.List[str], str, str, str, str, str) -> None
+@click.option(
+    "--var-encoding-json",
+    "-v",
+    default=None,
+    help=(
+        "encoding options to apply to all data variables in the dataset. "
+        "Can either be a JSON format string or the path to JSON file. "
+    ),
+)
+def to_netcdf(
+    inpaths, outpath, cdm, engine, backend_kwargs_json, netcdf_kwargs_json, var_encoding_json
+):  # type: (T.List[str], str, str, str, str, str, str) -> None
     import xarray as xr
 
     import cf2cdm
@@ -125,6 +134,12 @@ def to_netcdf(inpaths, outpath, cdm, engine, backend_kwargs_json, netcdf_kwargs_
     else:
         netcdf_kwargs = {}
 
+    if var_encoding_json is not None:
+        var_encoding = handle_json(var_encoding_json)
+        netcdf_kwargs.setdefault("encoding", {})
+        for var in ds.data_vars:
+            netcdf_kwargs["encoding"].setdefault(var, var_encoding)
+
     ds.to_netcdf(outpath, **netcdf_kwargs)
 
 
diff --git a/cfgrib/xarray_plugin.py b/cfgrib/xarray_plugin.py
index 1b6e5c82..ebbfce4d 100644
--- a/cfgrib/xarray_plugin.py
+++ b/cfgrib/xarray_plugin.py
@@ -105,7 +105,6 @@ def open_dataset(
         errors: str = "warn",
         extra_coords: T.Dict[str, str] = {},
     ) -> xr.Dataset:
-
         store = CfGribDataStore(
             filename_or_obj,
             indexpath=indexpath,
diff --git a/tests/test_60_main_commands.py b/tests/test_60_main_commands.py
index b202ab88..aa62fff7 100644
--- a/tests/test_60_main_commands.py
+++ b/tests/test_60_main_commands.py
@@ -73,6 +73,26 @@ def test_cfgrib_cli_to_netcdf_netcdf_kwargs(tmpdir: py.path.local) -> None:
     assert res.output == ""
 
 
+def test_cfgrib_cli_to_netcdf_var_encoding(tmpdir: py.path.local) -> None:
+    runner = click.testing.CliRunner()
+
+    var_encoding = '{"dtype": "float", "scale_factor": 0.1}'
+    res = runner.invoke(__main__.cfgrib_cli, ["to_netcdf", TEST_DATA, "-v", var_encoding])
+
+    assert res.exit_code == 0
+    assert res.output == ""
+
+    var_encoding_json = tmpdir.join("temp.json")
+    with open(var_encoding_json, "w") as f:
+        f.write(var_encoding)
+    res = runner.invoke(
+        __main__.cfgrib_cli, ["to_netcdf", TEST_DATA, "-v", str(var_encoding_json)]
+    )
+
+    assert res.exit_code == 0
+    assert res.output == ""
+
+
 def test_cfgrib_cli_dump() -> None:
     runner = click.testing.CliRunner()