Merge pull request #239 from ecmwf/feature/context_change

Feature/context change

polytope_feature/datacube/backends/datacube.py

@@ -148,17 +148,19 @@ def remap_path(self, path: DatacubePath):
         return path
 
     @staticmethod
-    def create(datacube, config={}, axis_options={}, compressed_axes_options=[], alternative_axes=[]):
+    def create(datacube, config={}, axis_options={}, compressed_axes_options=[], alternative_axes=[], context=None):
         # TODO: get the configs as None for pre-determined value and change them to empty dictionary inside the function
         if type(datacube).__name__ == "DataArray":
             from .xarray import XArrayDatacube
 
-            xadatacube = XArrayDatacube(datacube, axis_options, compressed_axes_options)
+            xadatacube = XArrayDatacube(datacube, axis_options, compressed_axes_options, context)
             return xadatacube
         if type(datacube).__name__ == "GribJump":
             from .fdb import FDBDatacube
 
-            fdbdatacube = FDBDatacube(datacube, config, axis_options, compressed_axes_options, alternative_axes)
+            fdbdatacube = FDBDatacube(
+                datacube, config, axis_options, compressed_axes_options, alternative_axes, context
+            )
             return fdbdatacube
 
     def check_branching_axes(self, request):

polytope_feature/datacube/backends/fdb.py

@@ -9,9 +9,13 @@
 
 
 class FDBDatacube(Datacube):
-    def __init__(self, gj, config=None, axis_options=None, compressed_axes_options=[], alternative_axes=[]):
+    def __init__(
+        self, gj, config=None, axis_options=None, compressed_axes_options=[], alternative_axes=[], context=None
+    ):
         if config is None:
             config = {}
+        if context is None:
+            context = {}
 
         super().__init__(axis_options, compressed_axes_options)
 
@@ -25,7 +29,7 @@ def __init__(self, gj, config=None, axis_options=None, compressed_axes_options=[
 
         self.gj = gj
         if len(alternative_axes) == 0:
-            self.fdb_coordinates = self.gj.axes(partial_request)
+            self.fdb_coordinates = self.gj.axes(partial_request, ctx=context)
             if len(self.fdb_coordinates) == 0:
                 raise BadRequestError(partial_request)
         else:

polytope_feature/datacube/backends/xarray.py

@@ -9,7 +9,7 @@
 class XArrayDatacube(Datacube):
     """Xarray arrays are labelled, axes can be defined as strings or integers (e.g. "time" or 0)."""
 
-    def __init__(self, dataarray: xr.DataArray, axis_options=None, compressed_axes_options=[]):
+    def __init__(self, dataarray: xr.DataArray, axis_options=None, compressed_axes_options=[], context=None):
         super().__init__(axis_options, compressed_axes_options)
         if axis_options is None:
             axis_options = {}

polytope_feature/polytope.py

@@ -39,7 +39,7 @@ def __repr__(self):
 
 
 class Polytope:
-    def __init__(self, datacube, engine=None, options=None):
+    def __init__(self, datacube, engine=None, options=None, context=None):
         from .datacube import Datacube
         from .engine import Engine
 
@@ -48,22 +48,24 @@ def __init__(self, datacube, engine=None, options=None):
 
         axis_options, compressed_axes_options, config, alternative_axes = PolytopeOptions.get_polytope_options(options)
 
-        self.datacube = Datacube.create(datacube, config, axis_options, compressed_axes_options, alternative_axes)
+        self.context = context
+
+        self.datacube = Datacube.create(
+            datacube, config, axis_options, compressed_axes_options, alternative_axes, self.context
+        )
         self.engine = engine if engine is not None else Engine.default()
         self.time = 0
 
     def slice(self, polytopes: List[ConvexPolytope]):
         """Low-level API which takes a polytope geometry object and uses it to slice the datacube"""
         return self.engine.extract(self.datacube, polytopes)
 
-    def retrieve(self, request: Request, method="standard", context=None):
+    def retrieve(self, request: Request, method="standard"):
         """Higher-level API which takes a request and uses it to slice the datacube"""
-        if context is None:
-            context = {}
-        logging.info("Starting request for %s ", context)
+        logging.info("Starting request for %s ", self.context)
         self.datacube.check_branching_axes(request)
         request_tree = self.engine.extract(self.datacube, request.polytopes())
-        logging.info("Created request tree for %s ", context)
-        self.datacube.get(request_tree, context)
-        logging.info("Retrieved data for %s ", context)
+        logging.info("Created request tree for %s ", self.context)
+        self.datacube.get(request_tree, self.context)
+        logging.info("Retrieved data for %s ", self.context)
         return request_tree