Add more tables

examples/message_dump.rs

@@ -83,9 +83,10 @@ fn main() {
                 Ok(f) => format!("{} {}", f.0, f.1.unwrap_or(0.0)),
                 Err(_) => "--".into(),
             },
-            match m.1.forecast_date() {
-                Ok(d) => format!("{d}"),
-                Err(_) => "--".into(),
+            match (m.1.time_interval_end(), m.1.forecast_date()) {
+                (Ok(None), Ok(d)) => format!("{d}"),
+                (Ok(Some(d)), _) => format!("{d}"),
+                _ => "--".into(),
             },
             match m.1.product_template_id() {
                 Ok(p) => format!("{p}"),

python/gribberish/gribberish_backend.py

@@ -27,6 +27,7 @@ def open_dataset(
         only_variables=None,
         perserve_dims=None,
         filter_by_attrs=None,
+        filter_by_variable_attrs=None,
     ):
         with open(filename_or_obj, 'rb') as f:
             raw_data = f.read()
@@ -36,7 +37,8 @@ def open_dataset(
                 drop_variables=drop_variables, 
                 only_variables=only_variables, 
                 perserve_dims=perserve_dims, 
-                filter_by_attrs=filter_by_attrs
+                filter_by_attrs=filter_by_attrs, 
+                filter_by_variable_attrs=filter_by_variable_attrs,
             )
             coords = {k: (v['dims'], v['values'], v['attrs']) for (k, v) in dataset['coords'].items()}
             data_vars = {k: (v['dims'], GribberishBackendArray(filename_or_obj, array_metadata=v['values']) , v['attrs']) for (k, v) in dataset['data_vars'].items()}
@@ -48,7 +50,14 @@ def open_dataset(
                 attrs=attrs
             )
 
-    open_dataset_parameters = ["filename_or_obj", "drop_variables", "only_variables", "perserve_dims", "filter_by_attrs"]
+    open_dataset_parameters = [
+        "filename_or_obj",
+        "drop_variables", 
+        "only_variables", 
+        "perserve_dims", 
+        "filter_by_attrs", 
+        "filter_by_variable_attrs"
+    ]
 
     def guess_can_open(self, filename_or_obj):
         try:

python/gribberish/kerchunk/mapper.py

@@ -92,6 +92,7 @@ def scan_gribberish(
     only_variables=None,
     perserve_dims=None,
     filter_by_attrs=None,
+    filter_by_variable_attrs=None,
     use_cfgrib_codec=False,
 ):
     """
@@ -113,6 +114,10 @@ def scan_gribberish(
         If given, dont shrink down these dimensions when their size is 1
     filter_by_attrs: dict
         If given, only store variables that match these attributes
+    filter_by_variable_attrs: dict
+        If given, check the attributes for these variables and only store
+            variables that match the listed attributes. If this is defined, 
+            filter_by_attrs is ignored.
     use_cfgrib_codec: bool
         If True, use the builtin kerchunk cfgrib codec instead of the
             default gribberish codec
@@ -128,7 +133,7 @@ def scan_gribberish(
     with fsspec.open(url, "rb", **storage_options) as f:
         for offset, size, data in _split_file(f):
             try:
-                dataset = parse_grib_dataset(data, perserve_dims=perserve_dims, encode_coords=True, filter_by_attrs=filter_by_attrs)
+                dataset = parse_grib_dataset(data, perserve_dims=perserve_dims, encode_coords=True, filter_by_attrs=filter_by_attrs, filter_by_variable_attrs=filter_by_variable_attrs)
                 var_name, var_data = next(iter(dataset['data_vars'].items()))
             except Exception as e:
                 # Skip messages that gribberish cannot handle yet or that are filtered out

python/src/dataset.rs

@@ -55,6 +55,7 @@ pub fn parse_grib_dataset<'py>(
     only_variables: Option<&PyList>,
     perserve_dims: Option<&PyList>,
     filter_by_attrs: Option<&PyDict>,
+    filter_by_variable_attrs: Option<&PyDict>,
     encode_coords: Option<bool>,
 ) -> PyResult<&'py PyDict> {
     let drop_variables = if let Some(drop_variables) = drop_variables {
@@ -92,6 +93,13 @@ pub fn parse_grib_dataset<'py>(
         PyDict::new(py)
     };
 
+    let (filter_by_variable_attrs, filter_by_variable_attrs_defined) =
+        if let Some(filter_by_variable_attrs) = filter_by_variable_attrs {
+            (filter_by_variable_attrs, true)
+        } else {
+            (PyDict::new(py), false)
+        };
+
     let encode_coords = if let Some(encode_coords) = encode_coords {
         encode_coords
     } else {
@@ -199,7 +207,11 @@ pub fn parse_grib_dataset<'py>(
     for (var, v) in var_mapping.iter() {
         let mut times = HashSet::new();
         for k in v.iter() {
-            times.insert(mapping.get(k).unwrap().2.forecast_date.clone());
+            if let Some(time_interval_end) = mapping.get(k).unwrap().2.time_interval_end {
+                times.insert(time_interval_end);
+            } else {
+                times.insert(mapping.get(k).unwrap().2.forecast_date.clone());
+            }
         }
         let mut times = times.into_iter().collect::<Vec<_>>();
         times.sort();
@@ -538,6 +550,15 @@ pub fn parse_grib_dataset<'py>(
         var_metadata
             .set_item("forecast_date", first.2.forecast_date.to_rfc3339())
             .unwrap();
+        var_metadata
+            .set_item(
+                "time_interval_end",
+                first
+                    .2
+                    .time_interval_end
+                    .map_or("".to_string(), |d| d.to_rfc3339()),
+            )
+            .unwrap();
         var_metadata
             .set_item(
                 "fixed_surface_type",
@@ -580,13 +601,32 @@ pub fn parse_grib_dataset<'py>(
         var_metadata.set_item("crs", first.2.proj.clone()).unwrap();
 
         let mut filtered = false;
-        for attr in var_metadata.keys() {
-            if filter_by_attrs.contains(attr).unwrap() {
-                let filter_value = filter_by_attrs.get_item(attr).unwrap().to_string();
-                let var_value = var_metadata.get_item(attr).unwrap().to_string();
-                if filter_value != var_value {
-                    filtered = true;
-                    break;
+        if filter_by_variable_attrs_defined {
+            if let Some(filter_by_variable_attrs) = filter_by_variable_attrs.get_item(var) {
+                let filter_by_variable_attrs = filter_by_variable_attrs
+                    .downcast::<PyDict>()
+                    .unwrap();
+                for attr in filter_by_variable_attrs.keys() {
+                    if filter_by_variable_attrs.contains(attr).unwrap() {
+                        let filter_value =
+                            filter_by_variable_attrs.get_item(attr).unwrap().to_string();
+                        let var_value = var_metadata.get_item(attr).unwrap().to_string();
+                        if filter_value != var_value {
+                            filtered = true;
+                            break;
+                        }
+                    }
+                }
+            }
+        } else {
+            for attr in var_metadata.keys() {
+                if filter_by_attrs.contains(attr).unwrap() {
+                    let filter_value = filter_by_attrs.get_item(attr).unwrap().to_string();
+                    let var_value = var_metadata.get_item(attr).unwrap().to_string();
+                    if filter_value != var_value {
+                        filtered = true;
+                        break;
+                    }
                 }
             }
         }

python/src/message.rs

@@ -60,14 +60,25 @@ impl GribMessageMetadata {
         self.inner.first_fixed_surface_value
     }
 
+    #[getter]
+    fn reference_date<'py>(&self, py: Python<'py>) -> PyResult<&'py PyDateTime> {
+        PyDateTime::from_timestamp(py, self.inner.reference_date.timestamp() as f64, None)
+    }
+
     #[getter]
     fn forecast_date<'py>(&self, py: Python<'py>) -> PyResult<&'py PyDateTime> {
         PyDateTime::from_timestamp(py, self.inner.forecast_date.timestamp() as f64, None)
     }
 
     #[getter]
-    fn reference_date<'py>(&self, py: Python<'py>) -> PyResult<&'py PyDateTime> {
-        PyDateTime::from_timestamp(py, self.inner.reference_date.timestamp() as f64, None)
+    fn time_interval_end<'py>(&self, py: Python<'py>) -> PyResult<Option<&'py PyDateTime>> {
+        if let Some(time_interval_end) = self.inner.time_interval_end {
+            let timestamp =
+                PyDateTime::from_timestamp(py, time_interval_end.timestamp() as f64, None)?;
+            Ok(Some(timestamp))
+        } else {
+            Ok(None)
+        }
     }
 
     #[getter]
@@ -158,11 +169,7 @@ impl GribMessage {
 }
 
 #[pyfunction]
-pub fn parse_grib_array<'py>(
-    py: Python<'py>,
-    data: &[u8],
-    offset: usize,
-) -> &'py PyArray1<f64> {
+pub fn parse_grib_array<'py>(py: Python<'py>, data: &[u8], offset: usize) -> &'py PyArray1<f64> {
     let message = Message::from_data(data, offset).unwrap();
 
     let mut data = message.data().unwrap();

src/message.rs

@@ -1,6 +1,6 @@
 use crate::sections::{indicator::Discipline, section::Section, section::SectionIterator};
 use crate::templates::product::tables::{
-    DerivedForecastType, FixedSurfaceType, GeneratingProcess, TypeOfStatisticalProcessing, TimeUnit,
+    DerivedForecastType, FixedSurfaceType, GeneratingProcess, TimeUnit, TypeOfStatisticalProcessing,
 };
 use crate::utils::iter::projection::LatLngProjection;
 use chrono::{DateTime, Utc};
@@ -425,10 +425,29 @@ impl<'a> Message<'a> {
             product_definition.product_definition_template(discipline as u8),
             "Only HorizontalAnalysisForecast templates are supported at this time".into()
         );
-        
+
         Ok(product_template.time_unit())
     }
 
+    pub fn time_interval_end(&self) -> Result<Option<DateTime<Utc>>, String> {
+        let discipline = self.discipline()?;
+
+        let product_definition = unwrap_or_return!(
+            self.sections().find_map(|s| match s {
+                Section::ProductDefinition(product_definition) => Some(product_definition),
+                _ => None,
+            }),
+            "Product definition section not found when reading variable data".into()
+        );
+
+        let product_template = unwrap_or_return!(
+            product_definition.product_definition_template(discipline as u8),
+            "Only HorizontalAnalysisForecast templates are supported at this time".into()
+        );
+
+        Ok(product_template.time_interval_end())
+    }
+
     pub fn first_fixed_surface(&self) -> Result<(FixedSurfaceType, Option<f64>), String> {
         let discipline = self.discipline()?;
 
@@ -548,10 +567,7 @@ impl<'a> Message<'a> {
             "Only latitude longitude templates supported at this time".into()
         );
 
-        Ok((
-            grid_template.y_count(),
-            grid_template.x_count(),
-        ))
+        Ok((grid_template.y_count(), grid_template.x_count()))
     }
 
     pub fn latlng_projector(&self) -> Result<LatLngProjection, String> {

src/message_metadata.rs

@@ -30,6 +30,7 @@ pub struct MessageMetadata {
     pub has_bitmap: bool,
     pub forecast_date: DateTime<Utc>,
     pub reference_date: DateTime<Utc>,
+    pub time_interval_end: Option<DateTime<Utc>>,
     pub proj: String,
     pub crs: String,
     pub is_regular_grid: bool,
@@ -79,8 +80,9 @@ impl<'a> TryFrom<&Message<'a>> for MessageMetadata {
                     .unwrap_or("Unknown".to_string())
             ),
             has_bitmap: message.has_bitmap(),
-            forecast_date: message.forecast_date()?,
             reference_date: message.reference_date()?,
+            forecast_date: message.forecast_date()?,
+            time_interval_end: message.time_interval_end()?,
             proj: message.proj_string()?,
             crs: message.crs()?,
             is_regular_grid: message.is_regular_grid()?,

src/templates/product/average_accumulation_extreme_horizontal_analysis_template.rs

@@ -171,7 +171,7 @@ impl ProductTemplate for AverageAccumulationExtremeHorizontalAnalysisForecastTem
         Some(self.data()[46].into())
     }
 
-    fn end_time(&self, _reference_date: DateTime<Utc>) -> Option<DateTime<Utc>> {
+    fn time_interval_end(&self) -> Option<DateTime<Utc>> {
         Some(self.time_interval_end())
     }
 }

src/templates/product/derived_ensemble_horizontal_analysis_template.rs

@@ -145,7 +145,7 @@ impl ProductTemplate for DerivedEnsembleHorizontalAnalysisForecastTemplate {
         None
     }
 
-    fn end_time(&self, _reference_date: DateTime<Utc>) -> Option<DateTime<Utc>> {
+    fn time_interval_end(&self) -> Option<DateTime<Utc>> {
         None
     }
 }

src/templates/product/derived_ensemble_horizontal_forecast_time_interval_template.rs

@@ -174,7 +174,7 @@ impl ProductTemplate for DerivedEnsembleHorizontalForecastTimeIntervalTemplate {
         Some(self.data()[48].into())
     }
 
-    fn end_time(&self, _reference_date: DateTime<Utc>) -> Option<DateTime<Utc>> {
+    fn time_interval_end(&self) -> Option<DateTime<Utc>> {
         Some(self.time_interval_end())
     }    
 }

src/templates/product/horizontal_analysis_template.rs

@@ -155,7 +155,7 @@ impl ProductTemplate for HorizontalAnalysisForecastTemplate {
         None
     }
 
-    fn end_time(&self, _reference_date: DateTime<Utc>) -> Option<DateTime<Utc>> {
+    fn time_interval_end(&self) -> Option<DateTime<Utc>> {
         None
     }
 }

src/templates/product/product_template.rs

@@ -10,7 +10,7 @@ pub trait ProductTemplate {
 	fn generating_process(&self) -> GeneratingProcess;
 	fn time_unit(&self) -> TimeUnit;
 	fn forecast_datetime(&self, reference_date: DateTime<Utc>) -> DateTime<Utc>;
-	fn end_time(&self, reference_date: DateTime<Utc>) -> Option<DateTime<Utc>>;
+	fn time_interval_end(&self) -> Option<DateTime<Utc>>;
     fn first_fixed_surface_type(&self) -> FixedSurfaceType;
 	fn first_fixed_surface_value(&self) -> Option<f64>;
     fn second_fixed_surface_type(&self) -> FixedSurfaceType;