//! 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, root_attrs: Arc>>, groups: Vec>>, } /// 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, 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, 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 { 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, cache_size: Option, headers: Option>, retries: Option, timeout: Option, allow_full_download: Option, max_full_download: Option, require_validator: Option, max_redirects: Option, max_parallel: Option, ) -> PyResult { 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>> { 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::("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) -> Py { 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 { 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> { 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>) -> PyResult> { self.read_file()?.get(py, key, default) } /// List the names of all children in the root group. fn keys(&self, py: Python<'_>) -> PyResult> { let names = self.read_file()?.member_names(py)?; Ok(PyList::new(py, names)?.into_any().unbind()) } fn values(&self, py: Python<'_>) -> PyResult> { let vals = self.read_file()?.values(py)?; Ok(PyList::new(py, vals)?.into_any().unbind()) } fn items(&self, py: Python<'_>) -> PyResult> { let items = self.read_file()?.items(py)?; Ok(PyList::new(py, items)?.into_any().unbind()) } fn __iter__(&self, py: Python<'_>) -> PyResult> { self.keys(py)?.call_method0(py, "__iter__") } fn __len__(&self, py: Python<'_>) -> PyResult { 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::( "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>, compression: Option<&str>, compression_opts: Option, ) -> 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> { 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 { 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::( "file is closed", )), } } fn __repr__(&self) -> String { match &self.inner { Some(FileInner::Read(root)) => match root.handle.redacted_url() { Some(url) => format!(""), None => format!("", self.filename), }, Some(FileInner::Write(s)) => { format!("", s.path.display()) } None => "".to_string(), } } fn __contains__(&self, py: Python<'_>, key: &str) -> PyResult { 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::( "cannot read from a file opened for writing", )), None => Err(PyErr::new::( "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::( "cannot write to a file opened for reading", )), None => Err(PyErr::new::( "file is closed", )), } } } fn parse_compression( compression: Option<&str>, compression_opts: Option, ) -> PyResult> { match compression { Some("gzip") => Ok(Some(compression_opts.unwrap_or(4))), Some(other) => Err(PyErr::new::(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(); } }