196 attempt to detect filename when saving

earthkit/data/core/fieldlist.py

@@ -13,7 +13,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
 
 
@@ -1164,16 +1164,17 @@ def _is_shared_grid(self):
                 )
         return False
 
+    @detect_out_filename
     def save(self, filename, append=False, **kwargs):
         r"""Write all the fields into a file.
 
         Parameters
         ----------
-        filename: str
-            The target file path.
-        append: bool
+        filename: str, optional
+            The target file path, if not defined attempts will be made to detect the filename
+        append: bool, optional
             When it is true append data to the target file. Otherwise
-            the target file be overwritten if already exists.
+            the target file be overwritten if already exists. Default is False
         **kwargs: dict, optional
             Other keyword arguments passed to :obj:`write`.
         """

earthkit/data/decorators.py

@@ -12,6 +12,7 @@
 import os
 import re
 import threading
+import warnings
 
 from earthkit.data.utils import load_json_or_yaml
 from earthkit.data.utils.availability import Availability
@@ -35,6 +36,38 @@ def wrapped(*args, **kwargs):
     return wrapped
 
 
+def detect_out_filename(func):
+    @functools.wraps(func)
+    def wrapped(self, *args, **kwargs):
+        # Detect filename:
+        if len(args) == 0:
+            for att in ["source_filename", "path"]:
+                if hasattr(self, att) and getattr(self, att) is not None:
+                    args = [os.path.basename(getattr(self, att))]
+                    break
+            else:
+                raise TypeError("Please provide an output filename")
+
+        # Ensure we do not overwrite file that is being read:
+        if (
+            args[0] is not None
+            and os.path.isfile(args[0])
+            and hasattr(self, "path")
+            and self.path is not None
+            and os.path.samefile(args[0], self.path)
+        ):
+            warnings.warn(
+                UserWarning(
+                    f"Earhtkit refusing to overwrite the file we are currently reading: {args[0]}"
+                )
+            )
+            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__)
 
@@ -32,6 +32,7 @@ def __init__(self, source, path):
 
         self._source = weakref.ref(source)
         self.path = path
+        self.source_filename = self.source.source_filename
 
     @property
     def source(self):
@@ -60,6 +61,7 @@ def ignore(self):
     def cache_file(self, *args, **kwargs):
         return self.source.cache_file(*args, **kwargs)
 
+    @detect_out_filename
     def save(self, path, **kwargs):
         mode = "wb" if self.binary else "w"
         with open(path, mode) as f:

earthkit/data/sources/__init__.py

@@ -33,6 +33,8 @@ class Source(Base):
     _dataset = None
     _parent = None
 
+    source_filename = None
+
     def __init__(self, **kwargs):
         self._kwargs = kwargs
 

earthkit/data/sources/cds.py

@@ -125,13 +125,16 @@ def __init__(self, dataset, *args, **kwargs):
 
     def _retrieve(self, dataset, request):
         def retrieve(target, args):
-            self.client().retrieve(args[0], args[1], target)
+            cds_result = self.client().retrieve(args[0], args[1])
+            self.source_filename = cds_result.location.split("/")[-1]
+            cds_result.download(target=target)
 
-        return self.cache_file(
+        return_object = self.cache_file(
             retrieve,
             (dataset, request),
             extension=EXTENSIONS.get(request.get("format"), ".cache"),
         )
+        return return_object
 
     @staticmethod
     @normalize("date", "date-list(%Y-%m-%d)")

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 earthkit.data.utils.parts import check_urls_and_parts, ensure_urls_and_parts
 
@@ -125,6 +126,7 @@ def to_numpy(self, **kwargs):
     def values(self):
         return self._reader.values
 
+    @detect_out_filename
     def save(self, path, **kwargs):
         return self._reader.save(path, **kwargs)
 

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
 
 
@@ -120,6 +123,42 @@ def test_cds_grib_multi_var_date(date, expected_date):
     assert s.metadata("date") == expected_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"))
+
+    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 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.source_filename))
+        os.chdir(here)
+
+
 @pytest.mark.long_test
 @pytest.mark.download
 @pytest.mark.skipif(NO_CDS, reason="No access to CDS")
@@ -212,6 +251,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

@@ -13,6 +13,8 @@
 import logging
 import os
 
+import pytest
+
 from earthkit.data import from_source
 from earthkit.data.core.temporary import temp_directory
 from earthkit.data.testing import earthkit_examples_file
@@ -28,11 +30,79 @@ def test_file_source_grib():
     assert len(s) == 2
 
 
+def test_file_source_grib_save():
+    s = from_source("file", earthkit_examples_file("test.grib"))
+    with temp_directory() as tmpdir:
+        # Check file save to assigned filename
+        s.save(os.path.join(tmpdir, "test2.grib"))
+        assert os.path.isfile(os.path.join(tmpdir, "test2.grib"))
+        # Check file can be saved in current dir with detected filename:
+        here = os.curdir
+        os.chdir(tmpdir)
+        s.save()
+        assert os.path.isfile("test.grib")
+        os.chdir(here)
+
+
+def test_file_source_grib_no_overwrite():
+    _s = from_source("file", earthkit_examples_file("test.grib"))
+    with temp_directory() as tmpdir:
+        os.chdir(tmpdir)
+        # Save the file locally
+        _s.save("test.grib")
+        # Open the local file
+        s = from_source("file", "test.grib")
+        with pytest.warns(
+            UserWarning,
+            match="Earhtkit refusing to overwrite the file we are currently reading",
+        ):
+            s.save("test.grib")
+        with pytest.warns(
+            UserWarning,
+            match="Earhtkit refusing to overwrite the file we are currently reading",
+        ):
+            s.save()
+
+
 def test_file_source_netcdf():
     s = from_source("file", earthkit_examples_file("test.nc"))
     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_netcdf_no_overwrite():
+    _s = from_source("file", earthkit_examples_file("test.nc"))
+    with temp_directory() as tmpdir:
+        os.chdir(tmpdir)
+        # Save the file locally
+        _s.save("test.nc")
+        # Open the local file
+        s = from_source("file", "test.nc")
+        with pytest.warns(
+            UserWarning,
+            match="Earhtkit refusing to overwrite the file we are currently reading",
+        ):
+            s.save("test.nc")
+        with pytest.warns(
+            UserWarning,
+            match="Earhtkit refusing to overwrite the file we are currently reading",
+        ):
+            s.save()
+
+
 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)