196 attempt to detect filename when saving

earthkit/data/core/fieldlist.py

@@ -12,7 +12,7 @@
 
 from earthkit.data.core import Base
 from earthkit.data.core.index import Index
-from earthkit.data.decorators import cached_method
+from earthkit.data.decorators import cached_method, detect_out_filename
 from earthkit.data.utils.metadata import metadata_argument
 
 
@@ -1153,6 +1153,7 @@ def _is_shared_grid(self):
                 )
         return False
 
+    @detect_out_filename
     def save(self, filename, append=False):
         r"""Write all the fields into a file.
 

earthkit/data/decorators.py

@@ -35,6 +35,26 @@ def wrapped(*args, **kwargs):
     return wrapped
 
 
+def detect_out_filename(func):
+    @functools.wraps(func)
+    def wrapped(self, *args, **kwargs):
+        if len(args) == 0:
+            try:
+                args = [os.path.basename(self.path)]
+            except AttributeError:
+                raise TypeError("Please provide a output filename")
+
+        if hasattr(self, "path") and args[0] == self.path:
+            LOG.warn(
+                "Earhtkit refusing to overwrite the file we are currently reading."
+            )
+            return
+
+        return func(self, *args, **kwargs)
+
+    return wrapped
+
+
 LOCK = threading.RLock()
 
 

earthkit/data/readers/__init__.py

@@ -14,7 +14,7 @@
 
 from earthkit.data.core import Base
 from earthkit.data.core.settings import SETTINGS
-from earthkit.data.decorators import locked
+from earthkit.data.decorators import detect_out_filename, locked
 
 LOG = logging.getLogger(__name__)
 
@@ -60,6 +60,7 @@ def ignore(self):
     def cache_file(self, *args, **kwargs):
         return self.source.cache_file(*args, **kwargs)
 
+    @detect_out_filename
     def save(self, path):
         mode = "wb" if self.binary else "w"
         with open(path, mode) as f:

earthkit/data/sources/file.py

@@ -14,6 +14,7 @@
 
 from earthkit.data import from_source
 from earthkit.data.core.caching import CACHE
+from earthkit.data.decorators import detect_out_filename
 from earthkit.data.readers import reader
 
 from . import Source
@@ -114,6 +115,7 @@ def to_numpy(self, **kwargs):
     def values(self):
         return self._reader.values
 
+    @detect_out_filename
     def save(self, path):
         return self._reader.save(path)
 

tests/sources/test_cds.py

@@ -9,9 +9,12 @@
 # nor does it submit to any jurisdiction.
 #
 
+import os
+
 import pytest
 
 from earthkit.data import from_source
+from earthkit.data.core.temporary import temp_directory
 from earthkit.data.testing import NO_CDS
 
 
@@ -116,6 +119,32 @@ def test_cds_grib_multi_var_date():
     ]
 
 
+@pytest.mark.long_test
+@pytest.mark.download
+@pytest.mark.skipif(NO_CDS, reason="No access to CDS")
+def test_cds_grib_save():
+    s = from_source(
+        "cds",
+        "reanalysis-era5-single-levels",
+        variable=["2t", "msl"],
+        product_type="reanalysis",
+        area=[50, -50, 20, 50],
+        date="2012-12-12",
+        time="12:00",
+    )
+    with temp_directory() as tmpdir:
+        # Check file save to assigned filename
+        s.save(os.path.join(tmpdir, "test.grib"))
+        assert os.path.isfile(os.path.join(tmpdir, "test.grib"))
+
+        # Check file can be saved in current dir with detected filename:
+        here = os.curdir
+        os.chdir(tmpdir)
+        s.save()
+        assert os.path.isfile(os.path.basename(s.path))
+        os.chdir(here)
+
+
 @pytest.mark.long_test
 @pytest.mark.download
 @pytest.mark.skipif(NO_CDS, reason="No access to CDS")
@@ -208,6 +237,33 @@ def test_cds_netcdf():
     assert s.metadata("variable") == ["t2m", "msl"]
 
 
+@pytest.mark.long_test
+@pytest.mark.download
+@pytest.mark.skipif(NO_CDS, reason="No access to CDS")
+def test_cds_netcdf_save():
+    s = from_source(
+        "cds",
+        "reanalysis-era5-single-levels",
+        variable=["2t", "msl"],
+        product_type="reanalysis",
+        area=[50, -50, 20, 50],
+        date="2012-12-12",
+        time="12:00",
+        format="netcdf",
+    )
+    with temp_directory() as tmpdir:
+        # Check file save to assigned filename
+        s.save(os.path.join(tmpdir, "test.nc"))
+        assert os.path.isfile(os.path.join(tmpdir, "test.nc"))
+
+        # Check file can be saved in current dir with detected filename:
+        here = os.curdir
+        os.chdir(tmpdir)
+        s.save()
+        assert os.path.isfile(os.path.basename(s.path))
+        os.chdir(here)
+
+
 @pytest.mark.long_test
 @pytest.mark.download
 @pytest.mark.skipif(NO_CDS, reason="No access to CDS")

tests/sources/test_file.py

@@ -33,6 +33,20 @@ def test_file_source_netcdf():
     assert len(s) == 2
 
 
+def test_file_source_netcdf_save():
+    s = from_source("file", earthkit_examples_file("test.nc"))
+    with temp_directory() as tmpdir:
+        # Check file save to assigned filename
+        s.save(os.path.join(tmpdir, "test2.nc"))
+        assert os.path.isfile(os.path.join(tmpdir, "test2.nc"))
+        # Check file can be saved in current dir with detected filename:
+        here = os.curdir
+        os.chdir(tmpdir)
+        s.save()
+        assert os.path.isfile("test.nc")
+        os.chdir(here)
+
+
 def test_file_source_odb():
     s = from_source("file", earthkit_examples_file("test.odb"))
     assert s.path == earthkit_examples_file("test.odb")

tests/translators/test_translators.py

@@ -18,6 +18,7 @@
 import xarray as xr
 
 from earthkit.data import from_source, transform, translators, wrappers
+from earthkit.data.testing import earthkit_test_data_file
 from earthkit.data.translators import ndarray as ndtranslator
 from earthkit.data.translators import pandas as pdtranslator
 from earthkit.data.translators import string as strtranslator
@@ -133,7 +134,7 @@ def test_gpd_dataframe_translator():
 
 def test_transform_from_grib_file():
     # transform grib-based data object
-    f = from_source("file", "tests/data/test_single.grib")
+    f = from_source("file", earthkit_test_data_file("test_single.grib"))
 
     # np.ndarray
     transformed = transform(f, np.ndarray)

tests/utils/test_module_inputs_wrapper.py

@@ -13,6 +13,7 @@
 
 from earthkit.data import from_object, from_source
 from earthkit.data.readers import Reader
+from earthkit.data.testing import earthkit_test_data_file
 from earthkit.data.utils import module_inputs_wrapper
 
 from . import dummy_module
@@ -58,7 +59,7 @@
 
 def test_transform_function_inputs_reader_to_xarray():
     # Check EK GribReader object
-    EK_GRIB_READER = from_source("file", "tests/data/test_single.grib")
+    EK_GRIB_READER = from_source("file", earthkit_test_data_file("test_single.grib"))
     ek_reader_result = WRAPPED_XR_ONES_LIKE(EK_GRIB_READER)
     # Will return a DataSet becuase that is first value in kwarg_types
     assert isinstance(ek_reader_result, xr.Dataset)
@@ -67,7 +68,7 @@ def test_transform_function_inputs_reader_to_xarray():
 
 def test_transform_function_inputs_reader_to_xarray_typesetting():
     # Check EK GribReader object
-    EK_GRIB_READER = from_source("file", "tests/data/test_single.grib")
+    EK_GRIB_READER = from_source("file", earthkit_test_data_file("test_single.grib"))
     ek_reader_result = WRAPPED_XR_ONES_LIKE_TYPE_SETTING(EK_GRIB_READER)
     # Will return a dataarray because that is first value in type-set Union
     assert isinstance(ek_reader_result, xr.DataArray)
@@ -76,7 +77,7 @@ def test_transform_function_inputs_reader_to_xarray_typesetting():
 
 def test_transform_module_inputs_reader_to_xarray():
     # Check EK GribReader object
-    EK_GRIB_READER = from_source("file", "tests/data/test_single.grib")
+    EK_GRIB_READER = from_source("file", earthkit_test_data_file("test_single.grib"))
     ek_reader_result = WRAPPED_DUMMY_MODULE.xarray_ones_like(EK_GRIB_READER)
     # Data array because type-setting of function has dataarray first
     assert isinstance(ek_reader_result, xr.DataArray)
@@ -108,22 +109,22 @@ def test_transform_module_inputs_wrapper_to_xarray():
 
 def test_transform_function_inputs_reader_to_numpy():
     # Test with Earthkit.data GribReader object
-    EK_GRIB_READER = from_source("file", "tests/data/test_single.grib")
+    EK_GRIB_READER = from_source("file", earthkit_test_data_file("test_single.grib"))
     assert WRAPPED_NP_MEAN(EK_GRIB_READER) == np.mean(EK_GRIB_READER.to_numpy())
     assert isinstance(WRAPPED_NP_MEAN(EK_GRIB_READER), np.float64)
 
 
 def test_transform_function_inputs_reader_to_numpy_typesetting():
     # Test with Earthkit.data GribReader object
-    EK_GRIB_READER = from_source("file", "tests/data/test_single.grib")
+    EK_GRIB_READER = from_source("file", earthkit_test_data_file("test_single.grib"))
     result = WRAPPED_NP_MEAN_TYPE_SETTING(EK_GRIB_READER)
     assert result == np.mean(EK_GRIB_READER.to_numpy())
     assert isinstance(result, np.float64)
 
 
 def test_transform_module_inputs_reader_to_numpy():
     # Test with Earthkit.data GribReader object
-    EK_GRIB_READER = from_source("file", "tests/data/test_single.grib")
+    EK_GRIB_READER = from_source("file", earthkit_test_data_file("test_single.grib"))
     result = WRAPPED_DUMMY_MODULE.numpy_mean(EK_GRIB_READER)
     assert result == np.mean(EK_GRIB_READER.to_numpy())
     assert isinstance(result, np.float64)