clawhdf5/crates/clawhdf5-py/src/attrs.rs

//! PyAttrs — dict-like access to HDF5 attributes.

use std::sync::{Arc, Mutex, PoisonError};

use clawhdf5_format::attribute::AttributeMessage;
use pyo3::exceptions::{PyKeyError, PyNotImplementedError, PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyList, PyTuple};

use crate::convert::{Converter, Elements, resolve_vl};
use crate::handle::Handle;
use crate::{OwnedAttrValue, PyEmpty, attr_value_to_py, edit, node, py_to_attr_value};

/// The attributes of an object in a file opened for reading (or editing).
struct ReadAttrs {
    handle: Arc<Handle>,
    addr: u64,
    path: String,
    /// Sorted by name, with the file generation they were read at: an edit
    /// (`attrs[name] = value`, here or through another handle on the same
    /// object) makes them re-read.
    cache: Mutex<(u64, Arc<Vec<AttributeMessage>>)>,
}

impl ReadAttrs {
    fn current(&self, py: Python<'_>) -> PyResult<Arc<Vec<AttributeMessage>>> {
        let generation = self.handle.generation();
        {
            let cached = self.cache.lock().unwrap_or_else(PoisonError::into_inner);
            if cached.0 == generation {
                return Ok(Arc::clone(&cached.1));
            }
        }
        let (addr, path) = (self.addr, &self.path);
        let attrs = Arc::new(self.handle.with(py, |f| node::attributes(f, addr, path))?);
        *self.cache.lock().unwrap_or_else(PoisonError::into_inner) =
            (generation, Arc::clone(&attrs));
        Ok(attrs)
    }

    fn check_writable(&self) -> PyResult<()> {
        if self.handle.is_writable() {
            return Ok(());
        }
        Err(PyErr::new::<pyo3::exceptions::PyIOError, _>(
            "cannot set attributes on a read-only file (open it with mode 'r+')",
        ))
    }

    fn set(&self, py: Python<'_>, name: &str, value: clawhdf5_rs::AttrValue) -> PyResult<()> {
        self.check_writable()?;
        let path = node::name(&self.path);
        self.handle.edit(py, |ed| ed.set_attr(&path, name, &value))
    }
}

/// Backing storage for attributes.
enum AttrsInner {
    Read(ReadAttrs),
    /// Writable attribute list shared with a parent (PyFile or PyGroup).
    Write(Arc<Mutex<Vec<(String, OwnedAttrValue)>>>),
}

/// Dict-like access to HDF5 attributes.
///
/// In read mode, values are what h5py returns: numpy scalars for scalar
/// attributes, numpy arrays otherwise, `str` for variable-length strings,
/// `numpy.bytes_` for fixed-length ones, and `Empty` for a null dataspace.
/// In a file opened with `'r+'`, `attrs[name] = value` adds or replaces an
/// attribute in the file at once (as h5py stores it, except that `str`
/// values become fixed-length UTF-8 strings). In write mode (`'w'`),
/// attributes set here are accumulated and written when the parent file is
/// closed.
#[pyclass(name = "Attrs")]
pub struct PyAttrs {
    inner: AttrsInner,
}

impl PyAttrs {
    /// The attributes of the object at `addr` (whose path is `path`) in a
    /// file opened for reading.
    pub(crate) fn read(
        py: Python<'_>,
        handle: Arc<Handle>,
        addr: u64,
        path: &str,
    ) -> PyResult<Self> {
        let generation = handle.generation();
        let attrs = Arc::new(handle.with(py, |f| node::attributes(f, addr, path))?);
        Ok(Self {
            inner: AttrsInner::Read(ReadAttrs {
                handle,
                addr,
                path: path.to_string(),
                cache: Mutex::new((generation, attrs)),
            }),
        })
    }

    /// Create a writable attrs that shares storage with a parent object.
    pub(crate) fn from_write(store: Arc<Mutex<Vec<(String, OwnedAttrValue)>>>) -> Self {
        Self {
            inner: AttrsInner::Write(store),
        }
    }

    fn set_value(
        &self,
        py: Python<'_>,
        key: &str,
        value: &Bound<'_, PyAny>,
        dtype: Option<&Bound<'_, PyAny>>,
        shape: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<()> {
        match &self.inner {
            AttrsInner::Read(r) => {
                r.check_writable()?;
                let value = edit::attr_value(py, value, dtype, shape)?;
                r.set(py, key, value)
            }
            AttrsInner::Write(store) => {
                if dtype.is_some() || shape.is_some() {
                    return Err(PyNotImplementedError::new_err(
                        "attrs.create with a dtype or shape is only supported in a file opened \
                         with 'r+'",
                    ));
                }
                let owned = py_to_attr_value(value)?;
                let mut guard = store.lock().unwrap();
                // Replace existing key if present.
                if let Some(entry) = guard.iter_mut().find(|(k, _)| k == key) {
                    entry.1 = owned;
                } else {
                    guard.push((key.to_string(), owned));
                }
                Ok(())
            }
        }
    }
}

#[pymethods]
impl PyAttrs {
    fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
        match &self.inner {
            AttrsInner::Read(r) => match r.current(py)?.iter().find(|a| a.name == key) {
                Some(attr) => Ok(attr_to_py(py, &r.handle, attr)?.unbind()),
                None => Err(PyKeyError::new_err(format!(
                    "Can't open attribute (can't locate attribute: '{key}')"
                ))),
            },
            AttrsInner::Write(store) => {
                let guard = store.lock().unwrap();
                for (k, v) in guard.iter() {
                    if k == key {
                        let attr_val: clawhdf5_rs::AttrValue = v.clone().into();
                        return Ok(attr_value_to_py(py, &attr_val));
                    }
                }
                Err(PyKeyError::new_err(key.to_string()))
            }
        }
    }

    /// `attrs[name] = value`. In a file opened with `'r+'` this writes the
    /// attribute (numeric, bool, complex, bytes and str data, any shape)
    /// into the file before returning; see the class docs.
    fn __setitem__(&self, py: Python<'_>, key: &str, value: &Bound<'_, PyAny>) -> PyResult<()> {
        self.set_value(py, key, value, None, None)
    }

    /// Deleting attributes is not supported in a file (the in-place editor
    /// cannot remove them); in write mode it removes a pending attribute.
    fn __delitem__(&self, key: &str) -> PyResult<()> {
        match &self.inner {
            AttrsInner::Read(_) => Err(PyNotImplementedError::new_err(format!(
                "cannot delete attribute '{key}': deleting attributes is not supported by \
                 clawhdf5's in-place editor"
            ))),
            AttrsInner::Write(store) => {
                let mut guard = store.lock().unwrap();
                let before = guard.len();
                guard.retain(|(k, _)| k != key);
                if guard.len() == before {
                    return Err(PyKeyError::new_err(key.to_string()));
                }
                Ok(())
            }
        }
    }

    /// h5py's `attrs.create(name, data, shape=None, dtype=None)`: `data`
    /// converted to `dtype` and reshaped to `shape` first.
    #[pyo3(signature = (name, data, shape=None, dtype=None))]
    fn create(
        &self,
        py: Python<'_>,
        name: &str,
        data: &Bound<'_, PyAny>,
        shape: Option<&Bound<'_, PyAny>>,
        dtype: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<()> {
        self.set_value(py, name, data, dtype, shape)
    }

    /// h5py's `attrs.modify(name, value)`: same as `attrs[name] = value`.
    fn modify(&self, py: Python<'_>, name: &str, value: &Bound<'_, PyAny>) -> PyResult<()> {
        self.set_value(py, name, value, None, None)
    }

    fn __len__(&self, py: Python<'_>) -> PyResult<usize> {
        match &self.inner {
            AttrsInner::Read(r) => Ok(r.current(py)?.len()),
            AttrsInner::Write(store) => Ok(store.lock().unwrap().len()),
        }
    }

    fn __contains__(&self, py: Python<'_>, key: &str) -> PyResult<bool> {
        match &self.inner {
            AttrsInner::Read(r) => Ok(r.current(py)?.iter().any(|a| a.name == key)),
            AttrsInner::Write(store) => Ok(store.lock().unwrap().iter().any(|(k, _)| k == key)),
        }
    }

    fn __iter__(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let keys = self.keys(py)?;
        let iter = keys.call_method0(py, "__iter__")?;
        Ok(iter)
    }

    fn __repr__(&self, py: Python<'_>) -> String {
        match self.__len__(py) {
            Ok(n) => format!("<HDF5 Attrs ({n} members)>"),
            Err(_) => "<HDF5 Attrs>".to_string(),
        }
    }

    /// The value of `key`, or `default` if there is no such attribute.
    #[pyo3(signature = (key, default=None))]
    fn get(&self, py: Python<'_>, key: &str, default: Option<Py<PyAny>>) -> PyResult<Py<PyAny>> {
        if self.__contains__(py, key)? {
            self.__getitem__(py, key)
        } else {
            Ok(default.unwrap_or_else(|| py.None()))
        }
    }

    /// Return attribute names as a list.
    fn keys(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let names: Vec<String> = match &self.inner {
            AttrsInner::Read(r) => r.current(py)?.iter().map(|a| a.name.clone()).collect(),
            AttrsInner::Write(store) => store
                .lock()
                .unwrap()
                .iter()
                .map(|(k, _)| k.clone())
                .collect(),
        };
        let list = PyList::new(py, &names)?;
        Ok(list.into_any().unbind())
    }

    /// Return attribute values as a list.
    fn values(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let vals: Vec<Py<PyAny>> = match &self.inner {
            AttrsInner::Read(r) => r
                .current(py)?
                .iter()
                .map(|a| attr_to_py(py, &r.handle, a).map(Bound::unbind))
                .collect::<PyResult<_>>()?,
            AttrsInner::Write(store) => store
                .lock()
                .unwrap()
                .iter()
                .map(|(_, v)| {
                    let attr: clawhdf5_rs::AttrValue = v.clone().into();
                    attr_value_to_py(py, &attr)
                })
                .collect(),
        };
        let list = PyList::new(py, &vals)?;
        Ok(list.into_any().unbind())
    }

    /// Return attribute (key, value) pairs as a list of tuples.
    fn items(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let pairs: Vec<(String, Py<PyAny>)> = match &self.inner {
            AttrsInner::Read(r) => r
                .current(py)?
                .iter()
                .map(|a| Ok((a.name.clone(), attr_to_py(py, &r.handle, a)?.unbind())))
                .collect::<PyResult<_>>()?,
            AttrsInner::Write(store) => store
                .lock()
                .unwrap()
                .iter()
                .map(|(k, v)| {
                    let attr: clawhdf5_rs::AttrValue = v.clone().into();
                    (k.clone(), attr_value_to_py(py, &attr))
                })
                .collect(),
        };
        let list = PyList::new(py, &pairs)?;
        Ok(list.into_any().unbind())
    }
}

/// An attribute's value as h5py returns it.
fn attr_to_py<'py>(
    py: Python<'py>,
    handle: &Handle,
    attr: &AttributeMessage,
) -> PyResult<Bound<'py, PyAny>> {
    crate::no_panic(|| {
        let conv = Converter::new(py, &attr.datatype, handle.offset_size)
            .map_err(|e| prefix_err(py, &attr.name, e))?;
        if node::is_null(&attr.dataspace) {
            return Ok(PyEmpty::new(conv.dtype).into_pyobject(py)?.into_any());
        }
        let shape: Vec<usize> = attr
            .dataspace
            .dimensions
            .iter()
            .map(|&d| d as usize)
            .collect();
        let n: usize = shape.iter().product();
        let data = if conv.is_vl() {
            let want = n * conv.elem_size;
            if attr.raw_data.len() < want {
                return Err(PyValueError::new_err(format!(
                    "attribute {}: {} bytes of variable-length references, expected {want}",
                    attr.name,
                    attr.raw_data.len(),
                )));
            }
            let raw = &attr.raw_data[..want];
            let (osz, lsz, unit) = (handle.offset_size, handle.length_size, conv.vl_unit);
            let what = format!("attribute {}", attr.name);
            Elements::Vl(handle.with(py, |f| {
                resolve_vl(f.storage(), raw, n, osz, lsz, unit).map_err(|e| e.into_py(&what))
            })?)
        } else {
            Elements::Bytes(attr.raw_data.clone())
        };
        let arr = conv
            .to_array(py, data, &shape, true)
            .map_err(|e| prefix_err(py, &attr.name, e))?;
        if shape.is_empty() {
            // A scalar dataspace: h5py returns the element itself.
            return arr.get_item(PyTuple::empty(py));
        }
        Ok(arr)
    })
}

fn prefix_err(py: Python<'_>, name: &str, e: PyErr) -> PyErr {
    let msg = format!("attribute {name}: {}", e.value(py));
    if e.is_instance_of::<PyTypeError>(py) {
        PyTypeError::new_err(msg)
    } else {
        PyValueError::new_err(msg)
    }
}