Add mars request dispatch

src/idpi/data_source.py

@@ -17,13 +17,6 @@
 # Local
 from . import config, mars
 
-GRIB_DEF = {
-    mars.Model.COSMO_1E: "cosmo",
-    mars.Model.COSMO_2E: "cosmo",
-    mars.Model.ICON_CH1_EPS: "cosmo",
-    mars.Model.ICON_CH2_EPS: "cosmo",
-}
-
 
 @contextmanager
 def cosmo_grib_defs():
@@ -98,7 +91,7 @@ def _(self, request: dict) -> Iterator:
         # validate the request
         req = mars.Request(**req_kwargs)
 
-        grib_def = config.get("data_scope", GRIB_DEF[req.model])
+        grib_def = config.get("data_scope", "cosmo")
         with grib_def_ctx(grib_def):
             if self.datafiles:
                 fs = ekd.from_source("file", self.datafiles)
@@ -115,12 +108,15 @@ def _(self, request: dict) -> Iterator:
                     asynchronous=False,
                 )
                 urls = [p["location"] for p in pointers]
-                print(urls)
                 source = ekd.from_source("url", urls)
             else:
                 source = ekd.from_source("fdb", req.to_fdb())
             yield from source  # type: ignore
 
+    @retrieve.register
+    def _(self, request: mars.Request) -> Iterator:
+        yield from self.retrieve(request.dump())
+
     @retrieve.register
     def _(self, request: str) -> Iterator:
         yield from self.retrieve({"param": request})

src/idpi/grib_decoder.py

@@ -17,7 +17,7 @@
 import xarray as xr
 
 # Local
-from . import data_source, metadata, tasking
+from . import data_source, mars, metadata, tasking
 
 logger = logging.getLogger(__name__)
 
@@ -28,7 +28,7 @@
 }
 INV_DIM_MAP = {v: k for k, v in DIM_MAP.items()}
 
-Request = str | tuple | dict
+Request = str | tuple | dict | mars.Request
 
 
 def _is_ensemble(field) -> bool:

tests/test_idpi/test_data_source.py

@@ -72,11 +72,10 @@ def test_retrieve_files_ifs(mock_from_source, mock_grib_def_ctx):
 
 
 def test_retrieve_fdb(mock_from_source, mock_grib_def_ctx):
-    datafiles = []
     param = "U"
     template = {"date": "20200101", "time": "0000"}
 
-    source = data_source.DataSource(datafiles, request_template=template)
+    source = data_source.DataSource(request_template=template)
     for _ in source.retrieve(param):
         pass
 
@@ -85,3 +84,19 @@ def test_retrieve_fdb(mock_from_source, mock_grib_def_ctx):
         call("fdb", mars.Request(param, **template).to_fdb()),
         call().__iter__(),
     ]
+
+
+def test_retrieve_fdb_mars(mock_from_source, mock_grib_def_ctx):
+    param = "U"
+    request = mars.Request(param=param)
+    template = {"date": "20200101", "time": "0000"}
+
+    source = data_source.DataSource(request_template=template)
+    for _ in source.retrieve(request):
+        pass
+
+    assert mock_grib_def_ctx.mock_calls == [call("cosmo")]
+    assert mock_from_source.mock_calls == [
+        call("fdb", mars.Request(param, **template).to_fdb()),
+        call().__iter__(),
+    ]