pygribjump: fix int pointer, add axes test

ecmwf · Nov 4, 2024 · 0229d90 · 0229d90
1 parent 89e89eb
commit 0229d90
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 5 deletions.
diff --git a/pygribjump/src/pygribjump/gribjump_c.h b/pygribjump/src/pygribjump/gribjump_c.h
@@ -26,7 +26,7 @@ int gribjump_result_values(gribjump_extraction_result_t* result, double*** value
 int gribjump_result_values_nocopy(gribjump_extraction_result_t* result, double*** values, unsigned long* nrange, unsigned long** nvalues);
 int gribjump_result_mask(gribjump_extraction_result_t* result,  unsigned long long*** masks, unsigned long* nrange, unsigned long** nmasks);
 int gribjump_delete_result(gribjump_extraction_result_t* result);
-int gribjump_new_axes(gj_axes_t** axes, const char* reqstr, int depth, const char* ctx, gribjump_handle_t* gj);
+int gribjump_new_axes(gj_axes_t** axes, const char* reqstr, int* level, const char* ctx, gribjump_handle_t* gj);
 int gribjump_axes_keys(gj_axes_t* axes, const char*** keys_out, unsigned long* size);
 int gribjump_axes_values(gj_axes_t* axes, const char* key, const char*** values_out, unsigned long* size);
 int gribjump_delete_axes(gj_axes_t* axes);

diff --git a/pygribjump/src/pygribjump/pygribjump.py b/pygribjump/src/pygribjump/pygribjump.py
@@ -224,7 +224,7 @@ def axes(self, req, level=3, ctx=None):
         requeststr = dic_to_request(req)
         newaxes = ffi.new('gj_axes_t**')
         reqstr = ffi.new('const char[]', requeststr.encode('ascii'))
-        level_c = ffi.new('int', level)
+        level_c = ffi.new('int*', level)
         lib.gribjump_new_axes(newaxes, reqstr, level_c, ctx_c, self.__gribjump)
 
         # TODO want to return a dict like:

diff --git a/pygribjump/tests/test_pygribjump.py b/pygribjump/tests/test_pygribjump.py
@@ -83,3 +83,16 @@ def test_extract_simple_sunshine_case(read_only_fdb_setup) -> None:
 
     actual = grib_jump.extract(polyrequest)
     assert numpy.array_equal(expected, actual[0][0][0][0], equal_nan=True)
+
+def test_axes(read_only_fdb_setup) -> None:
+    gribjump = pygj.GribJump()
+    req = {
+        "date": "20230508",
+    }
+    ax1 = gribjump.axes(req, level=1) # {'class': ['od'], 'date': ['20230508'], 'domain': ['g'], 'expver': ['0001'], 'stream': ['oper'], 'time': ['1200']}
+    ax2 = gribjump.axes(req, level=2) # {'class': ['od'], 'date': ['20230508'], 'domain': ['g'], 'expver': ['0001'], 'levtype': ['sfc'], 'stream': ['oper'], 'time': ['1200'], 'type': ['fc']}
+    ax3 = gribjump.axes(req, level=3) # {'class': ['od'], 'date': ['20230508'], 'domain': ['g'], 'expver': ['0001'], 'levelist': [''], 'levtype': ['sfc'], 'param': ['151130'], 'step': ['1'], 'stream': ['oper'], 'time': ['1200'], 'type': ['fc']}
+
+    assert len(ax1.keys()) == 6
+    assert len(ax2.keys()) == 8
+    assert len(ax3.keys()) == 11
diff --git a/src/gribjump/gribjump_c.cc b/src/gribjump/gribjump_c.cc
@@ -269,7 +269,7 @@ int extract(gribjump_handle_t* handle, gribjump_extraction_request_t** requests,
 }
 
 
-int gribjump_new_axes(gj_axes_t** axes, const char* reqstr, int depth, const char* ctx, gribjump_handle_t* gj) {
+int gribjump_new_axes(gj_axes_t** axes, const char* reqstr, int* level, const char* ctx, gribjump_handle_t* gj) {
     return wrapApiFunction([=] {
         ASSERT(gj);
         LogContext logctx;
@@ -278,7 +278,7 @@ int gribjump_new_axes(gj_axes_t** axes, const char* reqstr, int depth, const cha
         }
         std::string reqstr_str(reqstr);
         std::map<std::string, std::unordered_set<std::string>> values;
-        values = gj->axes(reqstr_str, depth, logctx);
+        values = gj->axes(reqstr_str, *level, logctx);
         *axes = new gj_axes_t(values);
     });
 }

diff --git a/src/gribjump/gribjump_c.h b/src/gribjump/gribjump_c.h
@@ -43,7 +43,7 @@ int gribjump_result_values_nocopy(gribjump_extraction_result_t* result, double**
 int gribjump_result_mask(gribjump_extraction_result_t* result,  unsigned long long*** masks, unsigned long* nrange, unsigned long** nmasks);
 int gribjump_delete_result(gribjump_extraction_result_t* result);
 
-int gribjump_new_axes(gj_axes_t** axes, const char* reqstr, int depth, const char* ctx, gribjump_handle_t* gj);
+int gribjump_new_axes(gj_axes_t** axes, const char* reqstr, int* level, const char* ctx, gribjump_handle_t* gj);
 int gribjump_axes_keys(gj_axes_t* axes, const char*** keys_out, unsigned long* size);
 int gribjump_axes_values(gj_axes_t* axes, const char* key, const char*** values_out, unsigned long* size);
 int gribjump_delete_axes(gj_axes_t* axes);