Feat/range types

design_notes.md

@@ -0,0 +1,31 @@
+# Questions
+
+* How are you supposed to interact with the Qube
+
+    * I think the compressed nodes and structure should be hidden from the user as much as possible
+
+    * Select/filter to return similar Qubes makes sense
+
+    * Would a "path" make sense?
+        qube[("class", &["od"])][("type", &["fc"])]
+        Not sure how it would work with coordinates. These don't point to a whole node. Would need to store coordinates along the path.
+        It would prevent nodes from being visible to the user.
+
+        is the path really any different to a filter though?
+        it stores the route separately to the result, which is quite nice perhaps
+
+        a path suggests you can only navigate down the tree, whereas a filter does not care about the structure of the tree
+        in other words a path reveals the hierarchy, whereas the qube does not really need to show that
+
+        A path does not make sense. Treat the Qube not as a tree from outside.
+
+    * Methods should include:
+        sel
+        dims()
+        coords(dim) -> list of coordinates along that dim
+
+        transpose (change internal ordering of axes)
+        align (align two Qubes to have the same axes?)
+        squeeze() removes length-1 dims.
+        drop_dim()
+        group_by() to group along a dimension by some function

py_qubed/src/lib.rs

@@ -1,16 +1,12 @@
 use ::qubed::Coordinates;
-use ::qubed::Datacube;
 use ::qubed::Qube;
 use ::qubed::select::SelectMode;
 use pyo3::exceptions::PyTypeError;
-use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyList, PyModule, PySet};
+use pyo3::types::{PyDict, PyList, PyModule};
 use serde_json::Value as JsonValue;
 
-mod helpers;
-
-#[pyclass(name = "Qube", unsendable)]
+#[pyclass(unsendable)]
 pub struct PyQube {
     inner: Qube,
 }
@@ -75,30 +71,6 @@ impl PyQube {
         }
     }
 
-    #[staticmethod]
-    #[pyo3(signature = (datacube, order=None))]
-    pub fn from_datacube(
-        datacube: Bound<'_, PyDict>,
-        order: Option<Vec<String>>,
-    ) -> PyResult<Self> {
-        let dc = pydict_to_datacube(datacube)?;
-        let order_slices: Option<&[String]> = order.as_deref();
-        Ok(Self { inner: Qube::from_datacube(&dc, order_slices) })
-    }
-
-    #[pyo3(signature = (datacube, order=None, accept_existing_order=false))]
-    pub fn append_datacube(
-        &mut self,
-        datacube: Bound<'_, PyDict>,
-        order: Option<Vec<String>>,
-        accept_existing_order: bool,
-    ) -> PyResult<()> {
-        let dc = pydict_to_datacube(datacube)?;
-        let order_slices: Option<&[String]> = order.as_deref();
-        self.inner.append_datacube(dc, order_slices, accept_existing_order);
-        Ok(())
-    }
-
     pub fn select(
         &self,
         request: Bound<'_, PyDict>,
@@ -200,7 +172,7 @@ impl PyQube {
         let mut validated_qubes = Vec::with_capacity(others.len());
         for item in others.iter() {
             let other_cell =
-                item.cast::<PyQube>().map_err(|_| PyTypeError::new_err("expected Qube"))?;
+                item.cast::<PyQube>().map_err(|_| PyTypeError::new_err("expected PyQube"))?;
             validated_qubes.push(other_cell.clone().unbind());
         }
 
@@ -216,290 +188,6 @@ impl PyQube {
     pub fn __repr__(&self) -> PyResult<String> {
         Ok(format!("PyQube(root_id={:?})", self.inner.root()))
     }
-
-    // -----------------------------------------------------------------------
-    //  New utility methods (mirrors of the @support_qubed_output functions
-    //  in forecast-in-a-box/fiab_plugin_ecmwf/qubed_utils.py)
-    // -----------------------------------------------------------------------
-
-    /// Return all dimension names and their union of coordinate values.
-    ///
-    /// This is an alias for `all_unique_dim_coords()` and returns the same
-    /// ``dict[str, list[str]]`` mapping.  The name `axes` matches the
-    /// terminology used in the qubed-utils helper library.
-    ///
-    /// ```python
-    /// q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"])
-    /// ax = q.axes()
-    /// assert set(ax["param"]) == {"2t", "tp"}
-    /// assert set(ax["time"])  == {"0", "1", "2"}
-    /// ```
-    pub fn axes(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
-        let dim_coords = self.inner.all_unique_dim_coords();
-        let py_dict = PyDict::new(py);
-
-        for (dimension, coordinates) in dim_coords {
-            let coord_str = coordinates.to_string();
-            let values: Vec<&str> = if coord_str.is_empty() {
-                vec![]
-            } else if coord_str.contains('/') {
-                coord_str.split('/').collect()
-            } else {
-                vec![&coord_str]
-            };
-
-            let py_list = PyList::empty(py);
-            for value in values {
-                py_list.append(value)?;
-            }
-            py_dict.set_item(dimension, py_list)?;
-        }
-
-        Ok(py_dict.into_any().unbind())
-    }
-
-    /// Return the set of all dimension names present anywhere in the Qube.
-    ///
-    /// ```python
-    /// q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"])
-    /// assert q.dimensions() == {"param", "time"}
-    /// ```
-    pub fn dimensions(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
-        let dims = self.inner.dimensions();
-        let py_set = PySet::new(py, dims.iter())?;
-        Ok(py_set.into_any().unbind())
-    }
-
-    /// Return the set of dimension names present in **every** leaf path.
-    ///
-    /// For a Qube with uniform depth this equals `dimensions()`.  When
-    /// branches differ in depth, only dimensions that appear in *all*
-    /// branches are returned.
-    ///
-    /// ```python
-    /// q1 = Qube.from_datacube({"param": "2t", "time": "0/1"}, ["param", "time"])
-    /// q2 = Qube.from_datacube({"param": "msl"}, ["param"])
-    /// q1.append(q2)
-    /// assert q1.common_dimensions() == {"param"}
-    /// ```
-    pub fn common_dimensions(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
-        let common = self.inner.common_dimensions();
-        let py_set = PySet::new(py, common.iter())?;
-        Ok(py_set.into_any().unbind())
-    }
-
-    /// Wrap the entire Qube under one or more new outer dimensions.
-    ///
-    /// `dimension` is a ``dict[str, list]`` where each key becomes a new
-    /// dimension name and the corresponding list supplies its coordinate
-    /// values.  Dimensions are applied in insertion order: the *last* entry
-    /// in the dict becomes the outermost dimension.
-    ///
-    /// The operation mutates the Qube in place and does **not** return a new
-    /// object (unlike the Python-reference implementation which follows an
-    /// immutable style).
-    ///
-    /// ```python
-    /// q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"])
-    /// q.expand({"ensemble": ["ens1", "ens2"]})
-    /// assert "ensemble" in q.dimensions()
-    /// assert "param"    in q.dimensions()
-    /// assert "time"     in q.dimensions()
-    /// ```
-    pub fn expand(&mut self, dimension: Bound<'_, PyDict>) -> PyResult<()> {
-        for (k, v) in dimension.iter() {
-            let key: String = k
-                .extract()
-                .map_err(|_| PyTypeError::new_err("expand: dimension keys must be strings"))?;
-
-            let coords = if v.is_instance_of::<PyList>() {
-                let lst =
-                    v.downcast::<PyList>().map_err(|e| PyTypeError::new_err(e.to_string()))?;
-                let joined = join_pylist_as_path(lst)?;
-                Coordinates::from_string(&joined)
-            } else {
-                let py_str = v.str()?;
-                let s: String = py_str.extract()?;
-                Coordinates::from_string(&s)
-            };
-
-            self.inner.expand(&key, coords).map_err(PyTypeError::new_err)?;
-        }
-        Ok(())
-    }
-
-    /// Remove one or more dimensions from the Qube.
-    ///
-    /// `axis` may be a single dimension name (``str``) or a list of names.
-    /// All children of the removed nodes are re-parented to their
-    /// grandparent, preserving the rest of the structure.  The result is
-    /// automatically compressed.
-    ///
-    /// Raises ``ValueError`` if any of the specified dimensions do not exist
-    /// in the Qube.
-    ///
-    /// ```python
-    /// q = Qube.from_datacube(
-    ///     {"param": "2t/tp", "time": "0/1/2", "level": "1000/850"},
-    ///     ["param", "time", "level"],
-    /// )
-    /// q.collapse("level")
-    /// assert "level" not in q.dimensions()
-    /// assert "param" in q.dimensions()
-    /// ```
-    pub fn collapse(&mut self, axis: Bound<'_, PyAny>) -> PyResult<()> {
-        // Accept either a single string or a list of strings.
-        let axes: Vec<String> = if axis.is_instance_of::<PyList>() {
-            let lst = axis.downcast::<PyList>().map_err(|e| PyTypeError::new_err(e.to_string()))?;
-            lst.iter()
-                .map(|item| {
-                    item.str()
-                        .and_then(|s| s.extract::<String>())
-                        .map_err(|_| PyTypeError::new_err("collapse: axis names must be strings"))
-                })
-                .collect::<PyResult<_>>()?
-        } else {
-            let s: String = axis
-                .extract()
-                .map_err(|_| PyTypeError::new_err("collapse: axis must be a str or list[str]"))?;
-            vec![s]
-        };
-
-        // Validate that every requested dimension actually exists.
-        let existing_dims = self.inner.dimensions();
-        let missing: Vec<&str> = axes
-            .iter()
-            .filter(|a| !existing_dims.contains(a.as_str()))
-            .map(String::as_str)
-            .collect();
-
-        if !missing.is_empty() {
-            return Err(PyValueError::new_err(format!(
-                "Dimension(s) '{}' not found in qube (existing: {:?})",
-                missing.join(", "),
-                {
-                    let mut sorted: Vec<&str> = existing_dims.iter().map(String::as_str).collect();
-                    sorted.sort();
-                    sorted
-                }
-            )));
-        }
-
-        self.inner.drop(axes).map_err(PyTypeError::new_err)
-    }
-
-    /// Collapse then expand in a single operation.
-    ///
-    /// Equivalent to calling `collapse(axis)` followed by `expand(dimension)`
-    /// on the same Qube.  Useful for replacing one dimension with another
-    /// (e.g. replacing "level" with "step").
-    ///
-    /// ```python
-    /// q = Qube.from_datacube(
-    ///     {"param": "2t/tp", "time": "0/1/2"},
-    ///     ["param", "time"],
-    /// )
-    /// q.coxpand("time", {"step": ["s1", "s2"]})
-    /// assert "time" not in q.dimensions()
-    /// assert "step" in q.dimensions()
-    /// ```
-    pub fn coxpand(
-        &mut self,
-        axis: Bound<'_, PyAny>,
-        dimension: Bound<'_, PyDict>,
-    ) -> PyResult<()> {
-        self.collapse(axis)?;
-        self.expand(dimension)
-    }
-
-    /// Check whether the Qube contains a given dimension or set of values.
-    ///
-    /// * **`str`** — returns ``True`` if the named dimension exists anywhere
-    ///   in the Qube.
-    /// * **`dict[str, list]`** — returns ``True`` if every key in the dict is
-    ///   a dimension that exists *and* every value listed for that key is
-    ///   present in that dimension's coordinate set.
-    /// * **`Qube`** — checks that every dimension/value present in the other
-    ///   Qube also exists in this one (subset check).
-    ///
-    /// ```python
-    /// q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"])
-    ///
-    /// assert q.contains("param")            is True
-    /// assert q.contains("level")            is False
-    /// assert q.contains({"param": ["2t"]})  is True
-    /// assert q.contains({"param": ["xyz"]}) is False
-    /// assert q.contains({"time": ["0", "1"]}) is True
-    /// ```
-    pub fn contains(&self, py: Python<'_>, item: Bound<'_, PyAny>) -> PyResult<bool> {
-        // Helper: convert a Coordinates value to a Vec<String> of individual values.
-        let coords_to_strings = |coords: &::qubed::Coordinates| -> Vec<String> {
-            let s = coords.to_string();
-            if s.is_empty() {
-                vec![]
-            } else if s.contains('/') {
-                s.split('/').map(|v| v.to_owned()).collect()
-            } else {
-                vec![s]
-            }
-        };
-
-        // Pre-compute current axes as Vec<String> per dimension.
-        let current_axes: std::collections::BTreeMap<String, Vec<String>> = self
-            .inner
-            .all_unique_dim_coords()
-            .into_iter()
-            .map(|(k, v)| (k, coords_to_strings(&v)))
-            .collect();
-
-        // --- branch 1: item is a plain string (dimension name check) ---
-        if let Ok(dim_str) = item.extract::<String>() {
-            return Ok(current_axes.contains_key(&dim_str));
-        }
-
-        // --- branch 2: item is a dict ---
-        if let Ok(dict) = item.downcast::<PyDict>() {
-            return helpers::check_dict_against_axes(dict, &current_axes, py);
-        }
-
-        // --- branch 3: item is a Qube ---
-        if let Ok(other_cell) = item.downcast::<PyQube>() {
-            let other_ref = other_cell.borrow();
-            let other_axes: std::collections::BTreeMap<String, Vec<String>> = other_ref
-                .inner
-                .all_unique_dim_coords()
-                .into_iter()
-                .map(|(k, v)| (k, coords_to_strings(&v)))
-                .collect();
-
-            for (key, other_vals) in &other_axes {
-                match current_axes.get(key) {
-                    None => return Ok(false),
-                    Some(cur_vals) => {
-                        for oval in other_vals {
-                            if !cur_vals.contains(oval) {
-                                return Ok(false);
-                            }
-                        }
-                    }
-                }
-            }
-            return Ok(true);
-        }
-
-        Err(PyTypeError::new_err("contains: item must be a str, dict[str, list], or Qube"))
-    }
-}
-
-fn pydict_to_datacube(datacube: Bound<'_, PyDict>) -> PyResult<Datacube> {
-    let mut dc = Datacube::new();
-    for (k, v) in datacube.iter() {
-        let key: String =
-            k.extract().map_err(|_| PyTypeError::new_err("datacube keys must be strings"))?;
-        let val_str: String = v.str()?.extract()?;
-        dc.add_coordinate(&key, Coordinates::from_string(&val_str));
-    }
-    Ok(dc)
 }
 
 pub(crate) fn join_pylist_as_path(lst: &Bound<'_, PyList>) -> PyResult<String> {