Python: Support obspec reader input to TIFF.open

python/Cargo.toml

@@ -19,6 +19,7 @@ crate-type = ["cdylib"]
 [dependencies]
 async-tiff = { path = "../" }
 bytes = "1.10.1"
+futures = "0.3.31"
 object_store = "0.12"
 pyo3 = { version = "0.23.0", features = ["macros"] }
 pyo3-async-runtimes = "0.23"

python/docs/api/tiff.md

@@ -3,3 +3,4 @@
 ::: async_tiff.TIFF
     options:
         show_if_no_docstring: true
+::: async_tiff.ObspecReader

python/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "maturin"
 [project]
 name = "async-tiff"
 requires-python = ">=3.9"
-dependencies = []
+dependencies = ["obspec>=0.1.0-beta.3"]
 dynamic = ["version"]
 classifiers = [
     "Programming Language :: Rust",

python/python/async_tiff/_async_tiff.pyi

@@ -3,5 +3,6 @@ from ._decoder import DecoderRegistry as DecoderRegistry
 from ._geo import GeoKeyDirectory as GeoKeyDirectory
 from ._ifd import ImageFileDirectory as ImageFileDirectory
 from ._thread_pool import ThreadPool as ThreadPool
+from ._tiff import ObspecInput as ObspecInput
 from ._tiff import TIFF as TIFF
 from ._tile import Tile as Tile

python/python/async_tiff/_tiff.pyi

@@ -1,15 +1,21 @@
-import obstore
+from typing import Protocol
 from ._tile import Tile
 from ._ifd import ImageFileDirectory
 from .store import ObjectStore
 
+# Fix exports
+from obspec._get import GetRangeAsync, GetRangesAsync
+
+class ObspecInput(GetRangeAsync, GetRangesAsync, Protocol):
+    """Supported obspec input to reader."""
+
 class TIFF:
     @classmethod
     async def open(
         cls,
         path: str,
         *,
-        store: obstore.store.ObjectStore | ObjectStore,
+        store: ObjectStore | ObspecInput,
         prefetch: int | None = 16384,
     ) -> TIFF:
         """Open a new TIFF.

python/src/lib.rs

@@ -4,6 +4,7 @@ mod decoder;
 mod enums;
 mod geo;
 mod ifd;
+mod reader;
 mod thread_pool;
 mod tiff;
 mod tile;

python/src/reader.rs

@@ -0,0 +1,128 @@
+use std::ops::Range;
+use std::sync::Arc;
+
+use async_tiff::error::{AsyncTiffError, AsyncTiffResult};
+use async_tiff::reader::{AsyncFileReader, ObjectReader};
+use bytes::Bytes;
+use futures::future::BoxFuture;
+use futures::FutureExt;
+use pyo3::exceptions::PyTypeError;
+use pyo3::intern;
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+use pyo3_async_runtimes::tokio::into_future;
+use pyo3_bytes::PyBytes;
+use pyo3_object_store::PyObjectStore;
+
+#[derive(FromPyObject)]
+pub(crate) enum StoreInput {
+    ObjectStore(PyObjectStore),
+    ObspecBackend(ObspecBackend),
+}
+
+impl StoreInput {
+    pub(crate) fn into_async_file_reader(self, path: String) -> Arc<dyn AsyncFileReader> {
+        match self {
+            Self::ObjectStore(store) => {
+                Arc::new(ObjectReader::new(store.into_inner(), path.into()))
+            }
+            Self::ObspecBackend(backend) => Arc::new(ObspecReader { backend, path }),
+        }
+    }
+}
+
+/// A Python backend for making requests that conforms to the GetRangeAsync and GetRangesAsync
+/// protocols defined by obspec.
+/// https://developmentseed.org/obspec/latest/api/get/#obspec.GetRangeAsync
+/// https://developmentseed.org/obspec/latest/api/get/#obspec.GetRangesAsync
+#[derive(Debug)]
+pub(crate) struct ObspecBackend(PyObject);
+
+impl ObspecBackend {
+    async fn get_range(&self, path: &str, range: Range<u64>) -> PyResult<PyBytes> {
+        let future = Python::with_gil(|py| {
+            let kwargs = PyDict::new(py);
+            kwargs.set_item(intern!(py, "path"), path)?;
+            kwargs.set_item(intern!(py, "start"), range.start)?;
+            kwargs.set_item(intern!(py, "end"), range.end)?;
+
+            let coroutine = self
+                .0
+                .call_method(py, intern!(py, "get_range"), (), Some(&kwargs))?;
+            into_future(coroutine.bind(py).clone())
+        })?;
+        let result = future.await?;
+        Python::with_gil(|py| result.extract(py))
+    }
+
+    async fn get_ranges(&self, path: &str, ranges: &[Range<u64>]) -> PyResult<Vec<PyBytes>> {
+        let starts = ranges.iter().map(|r| r.start).collect::<Vec<_>>();
+        let ends = ranges.iter().map(|r| r.end).collect::<Vec<_>>();
+
+        let future = Python::with_gil(|py| {
+            let kwargs = PyDict::new(py);
+            kwargs.set_item(intern!(py, "path"), path)?;
+            kwargs.set_item(intern!(py, "starts"), starts)?;
+            kwargs.set_item(intern!(py, "ends"), ends)?;
+
+            let coroutine = self
+                .0
+                .call_method(py, intern!(py, "get_range"), (), Some(&kwargs))?;
+            into_future(coroutine.bind(py).clone())
+        })?;
+        let result = future.await?;
+        Python::with_gil(|py| result.extract(py))
+    }
+
+    async fn get_range_wrapper(&self, path: &str, range: Range<u64>) -> AsyncTiffResult<Bytes> {
+        let result = self
+            .get_range(path, range)
+            .await
+            .map_err(|err| AsyncTiffError::External(Box::new(err)))?;
+        Ok(result.into_inner())
+    }
+
+    async fn get_ranges_wrapper(
+        &self,
+        path: &str,
+        ranges: Vec<Range<u64>>,
+    ) -> AsyncTiffResult<Vec<Bytes>> {
+        let result = self
+            .get_ranges(path, &ranges)
+            .await
+            .map_err(|err| AsyncTiffError::External(Box::new(err)))?;
+        Ok(result.into_iter().map(|b| b.into_inner()).collect())
+    }
+}
+
+impl<'py> FromPyObject<'py> for ObspecBackend {
+    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
+        let py = ob.py();
+        if ob.hasattr(intern!(py, "get_range_async"))?
+            && ob.hasattr(intern!(py, "get_ranges_async"))?
+        {
+            Ok(Self(ob.clone().unbind()))
+        } else {
+            Err(PyTypeError::new_err("Expected obspec-compatible class with `get_range_async` and `get_ranges_async` method."))
+        }
+    }
+}
+
+#[derive(Debug)]
+struct ObspecReader {
+    backend: ObspecBackend,
+    path: String,
+}
+
+impl AsyncFileReader for ObspecReader {
+    fn get_bytes(&self, range: Range<u64>) -> BoxFuture<'_, AsyncTiffResult<Bytes>> {
+        self.backend.get_range_wrapper(&self.path, range).boxed()
+    }
+
+    fn get_byte_ranges(
+        &self,
+        ranges: Vec<Range<u64>>,
+    ) -> BoxFuture<'_, AsyncTiffResult<Vec<Bytes>>> {
+        self.backend.get_ranges_wrapper(&self.path, ranges).boxed()
+    }
+}

python/src/tiff.rs

@@ -1,13 +1,13 @@
 use std::sync::Arc;
 
-use async_tiff::reader::{AsyncFileReader, ObjectReader, PrefetchReader};
+use async_tiff::reader::{AsyncFileReader, 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::AnyObjectStore;
 
+use crate::reader::StoreInput;
 use crate::tile::PyTile;
 use crate::PyImageFileDirectory;
 
@@ -25,25 +25,20 @@ impl PyTIFF {
         _cls: &'py Bound<PyType>,
         py: Python<'py>,
         path: String,
-        store: AnyObjectStore,
+        store: StoreInput,
         prefetch: Option<u64>,
     ) -> PyResult<Bound<'py, PyAny>> {
-        let reader = ObjectReader::new(store.into_dyn(), path.into());
-        let object_reader = reader.clone();
+        let reader = store.into_async_file_reader(path);
 
         let cog_reader = future_into_py(py, async move {
-            let reader: Box<dyn AsyncFileReader> = if let Some(prefetch) = prefetch {
-                Box::new(
-                    PrefetchReader::new(Box::new(reader), prefetch)
-                        .await
-                        .unwrap(),
-                )
+            let reader: Arc<dyn AsyncFileReader> = if let Some(prefetch) = prefetch {
+                Arc::new(PrefetchReader::new(reader, prefetch).await.unwrap())
             } else {
-                Box::new(reader)
+                reader
             };
             Ok(PyTIFF {
-                tiff: TIFF::try_open(reader).await.unwrap(),
-                reader: Arc::new(object_reader),
+                tiff: TIFF::try_open(reader.clone()).await.unwrap(),
+                reader,
             })
         })?;
         Ok(cog_reader)

src/cog.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use crate::error::AsyncTiffResult;
 use crate::ifd::ImageFileDirectories;
 use crate::reader::{AsyncCursor, AsyncFileReader};
@@ -13,7 +15,7 @@ impl TIFF {
     /// Open a new TIFF file.
     ///
     /// This will read all the Image File Directories (IFDs) in the file.
-    pub async fn try_open(reader: Box<dyn AsyncFileReader>) -> AsyncTiffResult<Self> {
+    pub async fn try_open(reader: Arc<dyn AsyncFileReader>) -> AsyncTiffResult<Self> {
         let mut cursor = AsyncCursor::try_open_tiff(reader).await?;
         let version = cursor.read_u16().await?;
 
@@ -72,13 +74,13 @@ mod test {
         let folder = "/Users/kyle/github/developmentseed/async-tiff/";
         let path = object_store::path::Path::parse("m_4007307_sw_18_060_20220803.tif").unwrap();
         let store = Arc::new(LocalFileSystem::new_with_prefix(folder).unwrap());
-        let reader = ObjectReader::new(store, path);
+        let reader = Arc::new(ObjectReader::new(store, path));
 
-        let cog_reader = TIFF::try_open(Box::new(reader.clone())).await.unwrap();
+        let cog_reader = TIFF::try_open(reader.clone()).await.unwrap();
 
         let ifd = &cog_reader.ifds.as_ref()[1];
         let decoder_registry = DecoderRegistry::default();
-        let tile = ifd.fetch_tile(0, 0, &reader).await.unwrap();
+        let tile = ifd.fetch_tile(0, 0, reader.as_ref()).await.unwrap();
         let tile = tile.decode(&decoder_registry).unwrap();
         std::fs::write("img.buf", tile).unwrap();
     }

src/error.rs

@@ -34,6 +34,10 @@ pub enum AsyncTiffError {
     /// Reqwest error
     #[error(transparent)]
     ReqwestError(#[from] reqwest::Error),
+
+    /// External error
+    #[error(transparent)]
+    External(Box<dyn std::error::Error + Send + Sync>),
 }
 
 /// Crate-specific result type.

src/reader.rs

@@ -165,13 +165,13 @@ impl AsyncFileReader for ReqwestReader {
 /// An AsyncFileReader that caches the first `prefetch` bytes of a file.
 #[derive(Debug)]
 pub struct PrefetchReader {
-    reader: Box<dyn AsyncFileReader>,
+    reader: Arc<dyn AsyncFileReader>,
     buffer: Bytes,
 }
 
 impl PrefetchReader {
     /// Construct a new PrefetchReader, catching the first `prefetch` bytes of the file.
-    pub async fn new(reader: Box<dyn AsyncFileReader>, prefetch: u64) -> AsyncTiffResult<Self> {
+    pub async fn new(reader: Arc<dyn AsyncFileReader>, prefetch: u64) -> AsyncTiffResult<Self> {
         let buffer = reader.get_bytes(0..prefetch).await?;
         Ok(Self { reader, buffer })
     }
@@ -213,14 +213,14 @@ pub(crate) enum Endianness {
 // TODO: in the future add buffering to this
 #[derive(Debug)]
 pub(crate) struct AsyncCursor {
-    reader: Box<dyn AsyncFileReader>,
+    reader: Arc<dyn AsyncFileReader>,
     offset: u64,
     endianness: Endianness,
 }
 
 impl AsyncCursor {
     /// Create a new AsyncCursor from a reader and endianness.
-    pub(crate) fn new(reader: Box<dyn AsyncFileReader>, endianness: Endianness) -> Self {
+    pub(crate) fn new(reader: Arc<dyn AsyncFileReader>, endianness: Endianness) -> Self {
         Self {
             reader,
             offset: 0,
@@ -230,7 +230,7 @@ impl AsyncCursor {
 
     /// Create a new AsyncCursor for a TIFF file, automatically inferring endianness from the first
     /// two bytes.
-    pub(crate) async fn try_open_tiff(reader: Box<dyn AsyncFileReader>) -> AsyncTiffResult<Self> {
+    pub(crate) async fn try_open_tiff(reader: Arc<dyn AsyncFileReader>) -> AsyncTiffResult<Self> {
         // Initialize with little endianness and then set later
         let mut cursor = Self::new(reader, Endianness::LittleEndian);
         let magic_bytes = cursor.read(2).await?;
@@ -252,7 +252,7 @@ impl AsyncCursor {
 
     /// Consume self and return the underlying [`AsyncFileReader`].
     #[allow(dead_code)]
-    pub(crate) fn into_inner(self) -> Box<dyn AsyncFileReader> {
+    pub(crate) fn into_inner(self) -> Arc<dyn AsyncFileReader> {
         self.reader
     }
 
@@ -316,7 +316,7 @@ impl AsyncCursor {
     }
 
     #[allow(dead_code)]
-    pub(crate) fn reader(&self) -> &dyn AsyncFileReader {
+    pub(crate) fn reader(&self) -> &Arc<dyn AsyncFileReader> {
         &self.reader
     }
 

tests/image_tiff/util.rs

@@ -10,6 +10,6 @@ const TEST_IMAGE_DIR: &str = "tests/image_tiff/images/";
 pub(crate) async fn open_tiff(filename: &str) -> TIFF {
     let store = Arc::new(LocalFileSystem::new_with_prefix(current_dir().unwrap()).unwrap());
     let path = format!("{TEST_IMAGE_DIR}/{filename}");
-    let reader = Box::new(ObjectReader::new(store.clone(), path.as_str().into()));
+    let reader = Arc::new(ObjectReader::new(store.clone(), path.as_str().into()));
     TIFF::try_open(reader).await.unwrap()
 }