Fetch tiles from Python (#48)

Author: kylebarron

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]: ...

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<dyn AsyncFileReader>,
+}
 
 #[pymethods]
 impl PyTIFF {
@@ -22,6 +29,7 @@ impl PyTIFF {
         prefetch: Option<u64>,
     ) -> PyResult<Bound<'py, PyAny>> {
         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<dyn AsyncFileReader> = 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<PyImageFileDirectory> {
-        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<Bound<'py, PyAny>> {
+        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<usize>,
+        y: Vec<usize>,
+        z: usize,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        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::<Vec<_>>();
+            Ok(py_tiles)
+        })
+    }
 }

python/src/tile.rs

@@ -13,6 +13,12 @@ use crate::PyDecoderRegistry;
 #[pyclass(name = "Tile")]
 pub(crate) struct PyTile(Option<Tile>);
 
+impl PyTile {
+    pub(crate) fn new(tile: Tile) -> Self {
+        Self(Some(tile))
+    }
+}
+
 #[pymethods]
 impl PyTile {
     #[getter]