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

//! PyFile — the main entry point for opening and creating HDF5 files.

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use pyo3::exceptions::{PyNotImplementedError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList};

use crate::attrs::PyAttrs;
use crate::group::{PyGroup, ReadGroup, WriteGroupState, finalize_write_group};
use crate::handle::Handle;
use crate::{DatasetSpec, OwnedAttrValue, apply_dataset_spec, extract_numpy_data, to_py_err};

/// Internal state for write mode.
struct WriteState {
    path: PathBuf,
    root_datasets: Vec<DatasetSpec>,
    root_attrs: Arc<Mutex<Vec<(String, OwnedAttrValue)>>>,
    groups: Vec<Arc<Mutex<WriteGroupState>>>,
}

/// An open HDF5 file.
///
/// Mirrors the h5py.File interface:
///
/// ```python
/// # Reading, a local file or a URL (range requests, nothing downloaded
/// # up front)
/// f = clawhdf5.File('data.h5', 'r')
/// f = clawhdf5.File('https://example.org/data.h5')
/// ds = f['dataset']
/// f.close()
///
/// # Writing
/// with clawhdf5.File('out.h5', 'w') as f:
///     f.create_dataset('data', data=numpy_array)
/// ```
#[pyclass(name = "File")]
pub struct PyFile {
    inner: Option<FileInner>,
    filename: String,
}

enum FileInner {
    /// The root group; it holds the file.
    Read(ReadGroup),
    Write(WriteState),
}

/// Whether `s` is a URL (`scheme://…`) rather than a path: the scheme is a
/// letter followed by letters, digits, `+`, `-` or `.` (RFC 3986).
fn is_url(s: &str) -> bool {
    let Some((scheme, _)) = s.split_once("://") else {
        return false;
    };
    let mut chars = scheme.chars();
    chars.next().is_some_and(|c| c.is_ascii_alphabetic())
        && chars.all(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '-' | '.'))
}

impl PyFile {
    fn from_handle(handle: Arc<Handle>, filename: String) -> Self {
        let root = handle.root;
        Self {
            inner: Some(FileInner::Read(ReadGroup::new(handle, String::new(), root))),
            filename,
        }
    }
}

#[pymethods]
impl PyFile {
    /// Open or create an HDF5 file.
    ///
    /// Parameters:
    ///   path: file path, or a URL (`http://`, `https://`, `s3://`, `gs://`,
    ///     `az://`; which schemes work depends on how the wheel was built)
    ///     to read the file remotely with default options (see `open_url`)
    ///   mode: 'r' for read (default), 'w' for write
    #[new]
    #[pyo3(signature = (path, mode="r"))]
    fn new(py: Python<'_>, path: &str, mode: &str) -> PyResult<Self> {
        let filename = path.to_string();
        if is_url(path) {
            if mode != "r" {
                return Err(PyValueError::new_err(format!(
                    "remote files are read-only: mode '{mode}' is not supported for a URL"
                )));
            }
            let handle = Handle::open_url(py, path, &clawhdf5_remote::Options::default())?;
            return Ok(Self::from_handle(handle, filename));
        }
        match mode {
            "r" => Ok(Self::from_handle(Handle::open_local(py, path)?, filename)),
            "r+" => Ok(Self::from_handle(
                Handle::open_editable(py, path)?,
                filename,
            )),
            "a" if std::path::Path::new(path).exists() => Ok(Self::from_handle(
                Handle::open_editable(py, path)?,
                filename,
            )),
            "a" => Err(PyNotImplementedError::new_err(format!(
                "mode 'a' on {path}, which does not exist: clawhdf5 can only edit an existing \
                 file in place; create a new one with mode 'w'"
            ))),
            "w" => Ok(Self {
                filename,
                inner: Some(FileInner::Write(WriteState {
                    path: PathBuf::from(path),
                    root_datasets: Vec::new(),
                    root_attrs: Arc::new(Mutex::new(Vec::new())),
                    groups: Vec::new(),
                })),
            }),
            other => Err(PyValueError::new_err(format!(
                "unsupported mode '{other}'; expected 'r', 'r+', 'a' or 'w'"
            ))),
        }
    }

    /// Open a remote file for reading, with options.
    ///
    /// The file is read through a block cache with range requests: opening
    /// costs one request (it also fetches the first block), and a read
    /// fetches only the blocks it needs. The GIL is released while waiting
    /// on the network.
    ///
    /// Parameters (all optional):
    ///   block_size: bytes per cached block (default 1 MiB)
    ///   cache_size: byte budget of the block cache (default 64 MiB)
    ///   headers: dict of extra HTTP headers (e.g. Authorization), sent only
    ///     to the URL's own origin
    ///   retries: retries of a request that failed transiently (default 3)
    ///   timeout: seconds to connect and receive response headers (default 30)
    ///   allow_full_download: when the server ignores Range requests,
    ///     download the whole file once instead of failing (default False)
    ///   max_full_download: largest file such a download may fetch
    ///     (default 1 GiB)
    ///   require_validator: refuse a server that sends neither ETag nor
    ///     Last-Modified (default False)
    ///   max_redirects: redirects followed per request (default 5)
    ///   max_parallel: requests of one read in flight at once (default 8)
    #[staticmethod]
    #[allow(clippy::too_many_arguments)]
    #[pyo3(signature = (url, *, block_size=None, cache_size=None, headers=None, retries=None,
        timeout=None, allow_full_download=None, max_full_download=None,
        require_validator=None, max_redirects=None, max_parallel=None))]
    fn open_url(
        py: Python<'_>,
        url: &str,
        block_size: Option<u64>,
        cache_size: Option<u64>,
        headers: Option<HashMap<String, String>>,
        retries: Option<u32>,
        timeout: Option<f64>,
        allow_full_download: Option<bool>,
        max_full_download: Option<u64>,
        require_validator: Option<bool>,
        max_redirects: Option<u32>,
        max_parallel: Option<usize>,
    ) -> PyResult<Self> {
        let mut options = clawhdf5_remote::Options::default();
        if let Some(b) = block_size {
            if b == 0 {
                return Err(PyValueError::new_err("block_size must be positive"));
            }
            options.cache.block_size = b;
            options.cache.coalesce_gap = b;
            // The opening request fetches the first block, not 1 MiB.
            options.http.first_request = b;
        }
        if let Some(c) = cache_size {
            options.cache.capacity = c;
        }
        let http = &mut options.http;
        if let Some(h) = headers {
            http.headers = h.into_iter().collect();
        }
        if let Some(r) = retries {
            http.retries = r;
        }
        if let Some(t) = timeout {
            if !(t.is_finite() && t > 0.0) {
                return Err(PyValueError::new_err("timeout must be a positive number"));
            }
            http.timeout = Duration::from_secs_f64(t);
        }
        if let Some(a) = allow_full_download {
            http.allow_full_download = a;
        }
        if let Some(m) = max_full_download {
            http.max_full_download = m;
        }
        if let Some(v) = require_validator {
            http.require_validator = v;
        }
        if let Some(r) = max_redirects {
            http.max_redirects = r;
        }
        if let Some(p) = max_parallel {
            if p == 0 {
                return Err(PyValueError::new_err("max_parallel must be positive"));
            }
            http.max_parallel = p;
        }
        let handle = Handle::open_url(py, url, &options)?;
        Ok(Self::from_handle(handle, url.to_string()))
    }

    /// For a remote file, what its block cache has done so far (reads,
    /// hits, misses, requests, bytes fetched, ...); `None` for a local file.
    #[getter]
    fn remote_stats<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyDict>>> {
        let Some(storage) = self.read_file()?.handle.remote_storage() else {
            return Ok(None);
        };
        let s = storage.stats();
        let d = PyDict::new(py);
        d.set_item("reads", s.reads)?;
        d.set_item("hits", s.hits)?;
        d.set_item("misses", s.misses)?;
        d.set_item("waits", s.waits)?;
        d.set_item("requests", s.requests)?;
        d.set_item("fetch_calls", s.fetch_calls)?;
        d.set_item("bytes_fetched", s.bytes_fetched)?;
        d.set_item("evictions", s.evictions)?;
        d.set_item("cached_bytes", s.cached_bytes)?;
        Ok(Some(d))
    }

    /// Close the file. In write mode, this finalizes and writes the file.
    fn close(&mut self) -> PyResult<()> {
        let inner = self.inner.take().ok_or_else(|| {
            PyErr::new::<pyo3::exceptions::PyIOError, _>("file is already closed")
        })?;
        match inner {
            FileInner::Read(root) => {
                root.handle.close();
                Ok(())
            }
            FileInner::Write(state) => finalize_write(state),
        }
    }

    /// Context manager entry — returns self.
    fn __enter__(slf: Py<Self>) -> Py<Self> {
        slf
    }

    /// Context manager exit — closes the file.
    #[pyo3(signature = (_exc_type=None, _exc_val=None, _exc_tb=None))]
    fn __exit__(
        &mut self,
        _exc_type: Option<&Bound<'_, PyAny>>,
        _exc_val: Option<&Bound<'_, PyAny>>,
        _exc_tb: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<bool> {
        self.close()?;
        Ok(false) // don't suppress exceptions
    }

    /// Get a child object (dataset or group) by path; `f['/']` is the root.
    fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
        self.read_file()?.get_item(py, key)
    }

    /// `f.get(key, default=None)`.
    #[pyo3(signature = (key, default=None))]
    fn get(&self, py: Python<'_>, key: &str, default: Option<Py<PyAny>>) -> PyResult<Py<PyAny>> {
        self.read_file()?.get(py, key, default)
    }

    /// List the names of all children in the root group.
    fn keys(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let names = self.read_file()?.member_names(py)?;
        Ok(PyList::new(py, names)?.into_any().unbind())
    }

    fn values(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let vals = self.read_file()?.values(py)?;
        Ok(PyList::new(py, vals)?.into_any().unbind())
    }

    fn items(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let items = self.read_file()?.items(py)?;
        Ok(PyList::new(py, items)?.into_any().unbind())
    }

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

    fn __len__(&self, py: Python<'_>) -> PyResult<usize> {
        Ok(self.read_file()?.member_names(py)?.len())
    }

    /// The root group's name, `/`.
    #[getter]
    fn name(&self) -> &'static str {
        "/"
    }

    /// `'r'` for a file opened read-only (a local file or a URL), `'r+'`
    /// for one open for editing or writing, as h5py reports it.
    #[getter]
    fn mode(&self) -> PyResult<&'static str> {
        match &self.inner {
            Some(FileInner::Read(root)) if !root.handle.is_writable() => Ok("r"),
            Some(_) => Ok("r+"),
            None => Err(PyErr::new::<pyo3::exceptions::PyIOError, _>(
                "file is closed",
            )),
        }
    }

    /// Nothing to do: every edit is written and synced when it is made, and
    /// a file opened with 'w' is written on `close()`.
    fn flush(&self) {}

    /// Deleting objects is not supported (h5py's `del f[name]`).
    fn __delitem__(&self, key: &str) -> PyResult<()> {
        Err(PyNotImplementedError::new_err(format!(
            "cannot delete '{key}': deleting objects is not supported by clawhdf5"
        )))
    }

    /// The path (or URL) the file was opened with.
    #[getter]
    fn filename(&self) -> &str {
        &self.filename
    }

    /// Create a dataset in the root group (write mode only).
    ///
    /// Parameters:
    ///   name: dataset name
    ///   data: numpy array
    ///   chunks: optional chunk dimensions (tuple or list)
    ///   compression: optional, only 'gzip' supported
    ///   compression_opts: gzip compression level (1-9)
    #[pyo3(signature = (name, *, data, chunks=None, compression=None, compression_opts=None))]
    fn create_dataset(
        &mut self,
        py: Python<'_>,
        name: &str,
        data: &Bound<'_, PyAny>,
        chunks: Option<Vec<u64>>,
        compression: Option<&str>,
        compression_opts: Option<u32>,
    ) -> PyResult<()> {
        let state = self.write_state_mut()?;
        let (dataset_data, shape) = extract_numpy_data(py, data)?;
        let deflate_level = parse_compression(compression, compression_opts)?;
        let spec = DatasetSpec {
            name: name.to_string(),
            data: dataset_data,
            shape,
            chunks,
            deflate_level,
            attrs: vec![],
        };
        state.root_datasets.push(spec);
        Ok(())
    }

    /// Create a group (write mode only). Returns a `Group` handle.
    fn create_group(&mut self, py: Python<'_>, name: &str) -> PyResult<Py<PyAny>> {
        let state = self.write_state_mut()?;
        let group_state = Arc::new(Mutex::new(WriteGroupState {
            name: name.to_string(),
            datasets: vec![],
            attrs: Arc::new(Mutex::new(vec![])),
        }));
        state.groups.push(Arc::clone(&group_state));
        let grp = PyGroup::from_write(group_state);
        Ok(grp.into_pyobject(py)?.into_any().unbind())
    }

    /// Attribute access. In read mode, returns attributes of the root group.
    /// In write mode, returns a writable attrs handle.
    #[getter]
    fn attrs(&self, py: Python<'_>) -> PyResult<PyAttrs> {
        match self.inner.as_ref() {
            Some(FileInner::Read(root)) => root.attrs(py),
            Some(FileInner::Write(state)) => Ok(PyAttrs::from_write(Arc::clone(&state.root_attrs))),
            None => Err(PyErr::new::<pyo3::exceptions::PyIOError, _>(
                "file is closed",
            )),
        }
    }

    fn __repr__(&self) -> String {
        match &self.inner {
            Some(FileInner::Read(root)) => match root.handle.redacted_url() {
                Some(url) => format!("<HDF5 File (read, \"{url}\")>"),
                None => format!("<HDF5 File (read, \"{}\")>", self.filename),
            },
            Some(FileInner::Write(s)) => {
                format!("<HDF5 File (write, \"{}\")>", s.path.display())
            }
            None => "<HDF5 File (closed)>".to_string(),
        }
    }

    fn __contains__(&self, py: Python<'_>, key: &str) -> PyResult<bool> {
        self.read_file()?.contains(py, key)
    }
}

impl PyFile {
    /// The root group of a file opened for reading.
    fn read_file(&self) -> PyResult<&ReadGroup> {
        match &self.inner {
            Some(FileInner::Read(f)) => Ok(f),
            Some(FileInner::Write(_)) => Err(PyErr::new::<pyo3::exceptions::PyIOError, _>(
                "cannot read from a file opened for writing",
            )),
            None => Err(PyErr::new::<pyo3::exceptions::PyIOError, _>(
                "file is closed",
            )),
        }
    }

    fn write_state_mut(&mut self) -> PyResult<&mut WriteState> {
        match &mut self.inner {
            Some(FileInner::Write(s)) => Ok(s),
            Some(FileInner::Read(root)) if root.handle.is_writable() => {
                Err(PyNotImplementedError::new_err(
                    "creating datasets or groups in an existing file is not supported by \
                     clawhdf5's in-place editor (mode 'r+' changes values, shapes and \
                     attributes)",
                ))
            }
            Some(FileInner::Read(_)) => Err(PyErr::new::<pyo3::exceptions::PyIOError, _>(
                "cannot write to a file opened for reading",
            )),
            None => Err(PyErr::new::<pyo3::exceptions::PyIOError, _>(
                "file is closed",
            )),
        }
    }
}

fn parse_compression(
    compression: Option<&str>,
    compression_opts: Option<u32>,
) -> PyResult<Option<u32>> {
    match compression {
        Some("gzip") => Ok(Some(compression_opts.unwrap_or(4))),
        Some(other) => Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
            "unsupported compression: {other}; only 'gzip' is supported"
        ))),
        None => Ok(None),
    }
}

/// Build and write the HDF5 file from accumulated write state.
fn finalize_write(state: WriteState) -> PyResult<()> {
    crate::no_panic(|| {
        let mut builder = clawhdf5_rs::FileBuilder::new();

        // Root attributes
        let root_attrs = state.root_attrs.lock().unwrap_or_else(|e| e.into_inner());
        for (name, val) in root_attrs.iter() {
            builder.set_attr(name, val.clone().into());
        }
        drop(root_attrs);

        // Root datasets
        for spec in &state.root_datasets {
            let db = builder.create_dataset(&spec.name);
            apply_dataset_spec(db, spec);
        }

        // Groups
        for group_arc in &state.groups {
            let guard = group_arc.lock().unwrap();
            finalize_write_group(&mut builder, &guard);
        }

        builder.write(&state.path).map_err(to_py_err)?;
        Ok(())
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn urls_and_paths() {
        assert!(is_url("http://h/f.h5"));
        assert!(is_url("s3://bucket/key.h5"));
        assert!(is_url("git+https://x"));
        assert!(!is_url("data.h5"));
        assert!(!is_url("/tmp/a://b.h5"));
        assert!(!is_url("dir/x://y"));
        assert!(!is_url("1http://x"));
        assert!(!is_url("://x"));
    }

    #[test]
    fn parse_gzip_compression() {
        assert_eq!(parse_compression(Some("gzip"), Some(6)).unwrap(), Some(6));
        assert_eq!(parse_compression(Some("gzip"), None).unwrap(), Some(4));
        assert_eq!(parse_compression(None, None).unwrap(), None);
        assert!(parse_compression(Some("lz4"), None).is_err());
    }

    #[test]
    fn finalize_roundtrip() {
        let dir = std::env::temp_dir();
        let path = dir.join("clawhdf5_py_test_finalize.h5");

        let state = WriteState {
            path: path.clone(),
            root_datasets: vec![DatasetSpec {
                name: "data".into(),
                data: crate::DatasetData::F64(vec![1.0, 2.0, 3.0]),
                shape: vec![3],
                chunks: None,
                deflate_level: None,
                attrs: vec![("unit".into(), OwnedAttrValue::Str("m".into()))],
            }],
            root_attrs: Arc::new(Mutex::new(vec![("version".into(), OwnedAttrValue::I64(1))])),
            groups: vec![Arc::new(Mutex::new(WriteGroupState {
                name: "grp".into(),
                datasets: vec![DatasetSpec {
                    name: "vals".into(),
                    data: crate::DatasetData::I32(vec![10, 20]),
                    shape: vec![2],
                    chunks: None,
                    deflate_level: None,
                    attrs: vec![],
                }],
                attrs: Arc::new(Mutex::new(vec![])),
            }))],
        };

        finalize_write(state).unwrap();

        let file = clawhdf5_rs::File::open(&path).unwrap();
        let ds = file.dataset("data").unwrap();
        assert_eq!(ds.read_f64().unwrap(), vec![1.0, 2.0, 3.0]);
        let grp_ds = file.dataset("grp/vals").unwrap();
        assert_eq!(grp_ds.read_i32().unwrap(), vec![10, 20]);

        std::fs::remove_file(&path).ok();
    }
}