diff --git a/python/python/async_tiff/_tiff.pyi b/python/python/async_tiff/_tiff.pyi index fb059f1..b5e1656 100644 --- a/python/python/async_tiff/_tiff.pyi +++ b/python/python/async_tiff/_tiff.pyi @@ -1,3 +1,4 @@ +from ._tile import Tile from ._ifd import ImageFileDirectory from .store import ObjectStore @@ -8,3 +9,5 @@ class TIFF: ) -> TIFF: ... @property def ifds(self) -> list[ImageFileDirectory]: ... + async def fetch_tile(self, x: int, y: int, z: int) -> Tile: ... + async def fetch_tiles(self, x: list[int], y: list[int], z: int) -> list[Tile]: ... diff --git a/python/src/tiff.rs b/python/src/tiff.rs index 1bf040e..27808e6 100644 --- a/python/src/tiff.rs +++ b/python/src/tiff.rs @@ -1,14 +1,21 @@ +use std::sync::Arc; + use async_tiff::reader::{AsyncFileReader, ObjectReader, PrefetchReader}; use async_tiff::TIFF; +use pyo3::exceptions::PyIndexError; use pyo3::prelude::*; use pyo3::types::PyType; use pyo3_async_runtimes::tokio::future_into_py; use pyo3_object_store::PyObjectStore; +use crate::tile::PyTile; use crate::PyImageFileDirectory; #[pyclass(name = "TIFF", frozen)] -pub(crate) struct PyTIFF(TIFF); +pub(crate) struct PyTIFF { + tiff: TIFF, + reader: Arc, +} #[pymethods] impl PyTIFF { @@ -22,6 +29,7 @@ impl PyTIFF { prefetch: Option, ) -> PyResult> { let reader = ObjectReader::new(store.into_inner(), path.into()); + let object_reader = reader.clone(); let cog_reader = future_into_py(py, async move { let reader: Box = if let Some(prefetch) = prefetch { @@ -33,14 +41,62 @@ impl PyTIFF { } else { Box::new(reader) }; - Ok(PyTIFF(TIFF::try_open(reader).await.unwrap())) + Ok(PyTIFF { + tiff: TIFF::try_open(reader).await.unwrap(), + reader: Arc::new(object_reader), + }) })?; Ok(cog_reader) } #[getter] fn ifds(&self) -> Vec { - let ifds = self.0.ifds(); + let ifds = self.tiff.ifds(); ifds.as_ref().iter().map(|ifd| ifd.clone().into()).collect() } + + fn fetch_tile<'py>( + &'py self, + py: Python<'py>, + x: usize, + y: usize, + z: usize, + ) -> PyResult> { + let reader = self.reader.clone(); + let ifd = self + .tiff + .ifds() + .as_ref() + .get(z) + .ok_or_else(|| PyIndexError::new_err(format!("No IFD found for z={}", z)))? + // TODO: avoid this clone; add Arc to underlying rust code? + .clone(); + future_into_py(py, async move { + let tile = ifd.fetch_tile(x, y, reader.as_ref()).await.unwrap(); + Ok(PyTile::new(tile)) + }) + } + + fn fetch_tiles<'py>( + &'py self, + py: Python<'py>, + x: Vec, + y: Vec, + z: usize, + ) -> PyResult> { + let reader = self.reader.clone(); + let ifd = self + .tiff + .ifds() + .as_ref() + .get(z) + .ok_or_else(|| PyIndexError::new_err(format!("No IFD found for z={}", z)))? + // TODO: avoid this clone; add Arc to underlying rust code? + .clone(); + future_into_py(py, async move { + let tiles = ifd.fetch_tiles(&x, &y, reader.as_ref()).await.unwrap(); + let py_tiles = tiles.into_iter().map(PyTile::new).collect::>(); + Ok(py_tiles) + }) + } } diff --git a/python/src/tile.rs b/python/src/tile.rs index 96b9b3d..feae546 100644 --- a/python/src/tile.rs +++ b/python/src/tile.rs @@ -13,6 +13,12 @@ use crate::PyDecoderRegistry; #[pyclass(name = "Tile")] pub(crate) struct PyTile(Option); +impl PyTile { + pub(crate) fn new(tile: Tile) -> Self { + Self(Some(tile)) + } +} + #[pymethods] impl PyTile { #[getter]