Proof of concept for opening COG in Python (#17)

Author: kylebarron

python/Cargo.toml

@@ -19,11 +19,25 @@ crate-type = ["cdylib"]
 [dependencies]
 async-tiff = { path = "../" }
 bytes = "1.8"
+# Match the version used by pyo3-object-store
+object_store = { git = "https://github.com/apache/arrow-rs", rev = "7a15e4b47ca97df2edef689c9f2ebd2f3888b79e" }
 pyo3 = { version = "0.23.0", features = ["macros"] }
+pyo3-async-runtimes = "0.23"
 pyo3-bytes = "0.1.2"
+pyo3-object_store = { git = "https://github.com/developmentseed/obstore", rev = "28ba07a621c1c104f084fb47ae7f8d08b1eae3ea" }
 thiserror = "1"
 tiff = "0.9.1"
 
+# We opt-in to using rustls as the TLS provider for reqwest, which is the HTTP
+# library used by object_store.
+# https://github.com/seanmonstar/reqwest/issues/2025
+reqwest = { version = "*", default-features = false, features = [
+    "rustls-tls-native-roots",
+] }
+
 [profile.release]
 lto = true
 codegen-units = 1
+
+[patch.crates-io]
+object_store = { git = "https://github.com/apache/arrow-rs", rev = "7a15e4b47ca97df2edef689c9f2ebd2f3888b79e" }

python/python/async_tiff/__init__.py

@@ -1,3 +1,4 @@
+from ._async_tiff import *
 from ._async_tiff import ___version
 
 __version__: str = ___version()

python/src/ifd.rs

@@ -215,3 +215,9 @@ impl PyImageFileDirectory {
         self.0.model_tiepoint()
     }
 }
+
+impl From<ImageFileDirectory> for PyImageFileDirectory {
+    fn from(value: ImageFileDirectory) -> Self {
+        Self(value)
+    }
+}

python/src/lib.rs

@@ -3,11 +3,13 @@
 mod enums;
 mod geo;
 mod ifd;
+mod tiff;
 
 use pyo3::prelude::*;
 
 use crate::geo::PyGeoKeyDirectory;
 use crate::ifd::PyImageFileDirectory;
+use crate::tiff::PyTIFF;
 
 const VERSION: &str = env!("CARGO_PKG_VERSION");
 
@@ -43,6 +45,10 @@ fn _async_tiff(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(___version))?;
     m.add_class::<PyGeoKeyDirectory>()?;
     m.add_class::<PyImageFileDirectory>()?;
+    m.add_class::<PyTIFF>()?;
+
+    pyo3_object_store::register_store_module(py, m, "async_tiff")?;
+    pyo3_object_store::register_exceptions_module(py, m, "async_tiff")?;
 
     Ok(())
 }

python/src/tiff.rs

@@ -0,0 +1,33 @@
+use async_tiff::{COGReader, ObjectReader};
+use pyo3::prelude::*;
+use pyo3::types::PyType;
+use pyo3_async_runtimes::tokio::future_into_py;
+use pyo3_object_store::PyObjectStore;
+
+use crate::PyImageFileDirectory;
+
+#[pyclass(name = "TIFF", frozen)]
+pub(crate) struct PyTIFF(COGReader);
+
+#[pymethods]
+impl PyTIFF {
+    #[classmethod]
+    #[pyo3(signature = (path, *, store))]
+    fn open<'py>(
+        _cls: &'py Bound<PyType>,
+        py: Python<'py>,
+        path: String,
+        store: PyObjectStore,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let reader = ObjectReader::new(store.into_inner(), path.into());
+        let cog_reader = future_into_py(py, async move {
+            Ok(PyTIFF(COGReader::try_open(Box::new(reader)).await.unwrap()))
+        })?;
+        Ok(cog_reader)
+    }
+
+    fn ifds(&self) -> Vec<PyImageFileDirectory> {
+        let ifds = self.0.ifds();
+        ifds.as_ref().iter().map(|ifd| ifd.clone().into()).collect()
+    }
+}

python/tests/__init__.py

@@ -0,0 +1,16 @@
+import async_tiff
+from async_tiff import TIFF
+from async_tiff.store import S3Store
+
+store = S3Store("sentinel-cogs", region="us-west-2", skip_signature=True)
+path = "sentinel-s2-l2a-cogs/12/S/UF/2022/6/S2B_12SUF_20220609_0_L2A/B04.tif"
+
+# 2 min, 15s
+tiff = await TIFF.open(path, store=store)
+ifds = tiff.ifds()
+ifd = ifds[0]
+ifd.tile_height
+ifd.tile_width
+ifd.photometric_interpretation
+gkd = ifd.geo_key_directory
+gkd.citation

python/tests/test_cog.py

@@ -1,4 +1,4 @@
-use std::io::{Cursor, SeekFrom};
+use std::io::Cursor;
 use std::ops::Range;
 use std::sync::Arc;
 
@@ -25,7 +25,7 @@ use crate::error::{AiocogeoError, Result};
 /// [`ObjectStore`]: object_store::ObjectStore
 ///
 /// [`tokio::fs::File`]: https://docs.rs/tokio/latest/tokio/fs/struct.File.html
-pub trait AsyncFileReader: Send {
+pub trait AsyncFileReader: Send + Sync {
     /// Retrieve the bytes in `range`
     fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>>;
 
@@ -57,12 +57,12 @@ impl AsyncFileReader for Box<dyn AsyncFileReader + '_> {
 }
 
 #[cfg(feature = "tokio")]
-impl<T: tokio::io::AsyncRead + tokio::io::AsyncSeek + Unpin + Send> AsyncFileReader for T {
+impl<T: tokio::io::AsyncRead + tokio::io::AsyncSeek + Unpin + Send + Sync> AsyncFileReader for T {
     fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
         use tokio::io::{AsyncReadExt, AsyncSeekExt};
 
         async move {
-            self.seek(SeekFrom::Start(range.start)).await?;
+            self.seek(std::io::SeekFrom::Start(range.start)).await?;
 
             let to_read = (range.end - range.start).try_into().unwrap();
             let mut buffer = Vec::with_capacity(to_read);

src/async_reader.rs

@@ -25,6 +25,10 @@ impl COGReader {
         Ok(Self { reader, ifds })
     }
 
+    pub fn ifds(&self) -> &ImageFileDirectories {
+        &self.ifds
+    }
+
     /// Return the EPSG code representing the crs of the image
     pub fn epsg(&self) -> Option<u16> {
         let ifd = &self.ifds.as_ref()[0];