Handle missing fields

Author: AreWeDreaming

imas/test/test_wrangle.py

@@ -153,7 +153,8 @@ def get_dtype(arr):
         return type(arr)
 
 def test_unwrangle(test_ids_dict, flat):
-    result = unwrangle(list(flat.keys()), test_ids_dict)
+    result, failed = unwrangle(list(flat.keys()), test_ids_dict)
+    assert failed == 0, f"The following fields failed to load {failed}"
     for key in flat.keys():
         if np.issubdtype(get_dtype(result[key]), np.floating):
             assert ak.almost_equal(result[key], flat[key])

imas/wrangler.py

@@ -56,6 +56,7 @@ def unwrangle(
 ) -> Dict[str, ak.Array | np.ndarray]:
     flat = {}  
     ids_locations = split_location_across_ids(locations)
+    failed_locations = []
     for key in ids_locations:
         tensorizer = IDSTensorizer(ids_dict[key], ids_locations[key])
         tensorizer.include_coordinate_paths()
@@ -64,7 +65,11 @@ def unwrangle(
         # Add IDS conversion
         for ids_location in ids_locations[key]:
             location = key + "." + ids_location.replace("/", ".")
-            values = tensorizer.awkward_tensorize(ids_location)
+            try:
+                values = tensorizer.awkward_tensorize(ids_location)
+            except KeyError:
+                failed_locations.append(location)
+                continue
             if hasattr(values, "__getitem__"):
                 # Not a scalar, e.g. homogenous_time
                 try:
@@ -73,4 +78,4 @@ def unwrangle(
                     flat[location] = ak.Array(values)
             else:
                 flat[location] = values
-    return flat
+    return flat, failed_locations