WIP: requester pays

pyo3-object_store/Cargo.toml

@@ -15,6 +15,7 @@ include = ["src", "type-hints", "README.md", "LICENSE"]
 
 [dependencies]
 futures = "0.3"
+http = "1.0.0"
 object_store = { version = "0.11", features = ["aws", "azure", "gcp", "http"] }
 pyo3 = { version = "0.22", features = ["chrono", "indexmap"] }
 pyo3-async-runtimes = { version = "0.22", features = ["tokio-runtime"] }

pyo3-object_store/src/api.rs

@@ -2,7 +2,9 @@ use pyo3::intern;
 use pyo3::prelude::*;
 
 use crate::error::*;
-use crate::{PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyS3Store};
+use crate::{
+    PyAzureStore, PyClientOptions, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyS3Store,
+};
 
 /// Export the default Python API as a submodule named `store` within the given parent module
 ///
@@ -50,6 +52,7 @@ pub fn register_store_module(
     child_module.add_class::<PyLocalStore>()?;
     child_module.add_class::<PyMemoryStore>()?;
     child_module.add_class::<PyS3Store>()?;
+    child_module.add_class::<PyClientOptions>()?;
 
     parent_module.add_submodule(&child_module)?;
 

pyo3-object_store/src/client.rs

@@ -1,9 +1,11 @@
 use std::collections::HashMap;
 use std::str::FromStr;
 
+use http::{HeaderMap, HeaderName, HeaderValue};
 use object_store::{ClientConfigKey, ClientOptions};
 use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedStr;
+use pyo3::types::PyDict;
 
 use crate::error::PyObjectStoreError;
 
@@ -37,10 +39,15 @@ impl<'py> FromPyObject<'py> for PyClientConfigValue {
 
 /// A wrapper around `ClientOptions` that implements [`FromPyObject`].
 #[derive(Debug)]
+#[pyclass(name = "ClientOptions", frozen)]
 pub struct PyClientOptions(ClientOptions);
 
 impl<'py> FromPyObject<'py> for PyClientOptions {
     fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
+        if let Ok(options) = ob.downcast::<PyClientOptions>() {
+            return Ok(Self(options.get().0.clone()));
+        }
+
         let py_input = ob.extract::<HashMap<PyClientConfigKey, String>>()?;
         let mut options = ClientOptions::new();
         for (key, value) in py_input.into_iter() {
@@ -55,3 +62,75 @@ impl From<PyClientOptions> for ClientOptions {
         value.0
     }
 }
+
+#[pymethods]
+impl PyClientOptions {
+    #[new]
+    #[pyo3(signature = (*, default_headers = None, **kwargs))]
+    // TODO: add kwargs
+    fn py_new(
+        default_headers: Option<PyHeaderMap>,
+        kwargs: Option<&Bound<PyDict>>,
+    ) -> PyResult<Self> {
+        let mut options = ClientOptions::default();
+        if let Some(default_headers) = default_headers {
+            options = options.with_default_headers(default_headers.0);
+        }
+        if let Some(kwargs) = kwargs {
+            let kwargs = kwargs.extract::<HashMap<PyClientConfigKey, String>>()?;
+            for (key, value) in kwargs.into_iter() {
+                options = options.with_config(key.0, value);
+            }
+        }
+
+        Ok(Self(options))
+    }
+}
+
+// use pyo3::prelude::*;
+// use pyo3::types::PyDict;
+
+// #[pyfunction]
+// #[pyo3(signature = (**kwds))]
+// fn num_kwds(kwds: Option<&Bound<'_, PyDict>>) -> usize {
+//     kwds.map_or(0, |dict| dict.len())
+// }
+
+// #[pymodule]
+// fn module_with_functions(m: &Bound<'_, PyModule>) -> PyResult<()> {
+//     m.add_function(wrap_pyfunction!(num_kwds, m)?)
+// }
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct PyHeaderName(HeaderName);
+
+impl<'py> FromPyObject<'py> for PyHeaderName {
+    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
+        // TODO: check that this works on both str and bytes input
+        Ok(Self(HeaderName::from_bytes(ob.extract()?).unwrap()))
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct PyHeaderValue(HeaderValue);
+
+impl<'py> FromPyObject<'py> for PyHeaderValue {
+    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
+        Ok(Self(
+            HeaderValue::from_str(ob.extract::<PyBackedStr>()?.as_ref()).unwrap(),
+        ))
+    }
+}
+
+pub struct PyHeaderMap(HeaderMap);
+
+impl<'py> FromPyObject<'py> for PyHeaderMap {
+    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
+        let py_input = ob.extract::<HashMap<PyHeaderName, PyHeaderValue>>()?;
+        let mut header = HeaderMap::with_capacity(py_input.len());
+        for (key, value) in py_input.into_iter() {
+            header.insert(key.0, value.0);
+        }
+        Ok(Self(header))
+    }
+}

pyo3-object_store/type-hints/__init__.pyi

@@ -6,6 +6,7 @@ from ._aws import S3Store as S3Store
 from ._azure import AzureConfigKey as AzureConfigKey
 from ._azure import AzureStore as AzureStore
 from ._client import ClientConfigKey as ClientConfigKey
+from ._client import ClientOptions as ClientOptions
 from ._gcs import GCSConfigKey as GCSConfigKey
 from ._gcs import GCSStore as GCSStore
 from ._http import HTTPStore as HTTPStore

pyo3-object_store/type-hints/_aws.pyi

@@ -4,7 +4,7 @@ import boto3
 import botocore
 import botocore.session
 
-from ._client import ClientConfigKey
+from ._client import ClientConfigKey, ClientOptions
 from ._retry import RetryConfig
 
 S3ConfigKey = Literal[
@@ -164,7 +164,7 @@ class S3Store:
         bucket: str,
         *,
         config: Dict[S3ConfigKey | str, str] | None = None,
-        client_options: Dict[ClientConfigKey, str | bool] | None = None,
+        client_options: Dict[ClientConfigKey, str | bool] | ClientOptions | None = None,
         retry_config: RetryConfig | None = None,
     ) -> S3Store:
         """Construct a new S3Store with credentials inferred from a boto3 Session

pyo3-object_store/type-hints/_client.pyi

@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Dict, Literal
 
 ClientConfigKey = Literal[
     "allow_http",
@@ -61,3 +61,6 @@ Either lower case or upper case strings are accepted.
     response body has finished.
 - `"user_agent"`: User-Agent header to be used by this client.
 """
+
+class ClientOptions:
+    def __init__(self, *, default_headers: Dict[str, str], **kwargs) -> None: ...

tests/test_get.py

@@ -1,7 +1,24 @@
 import pytest
 
 import obstore as obs
-from obstore.store import MemoryStore
+import obstore
+
+dir(obstore)
+import obstore.store
+
+dir(obstore.store)
+
+from obstore.store import MemoryStore, ClientOptions, S3Store
+
+import boto3
+
+session = boto3.Session()
+
+options = ClientOptions(default_headers={b"x-amz-request-payer": "requester"})
+store = S3Store.from_session(session, "naip-visualization", client_options=options)
+result = obs.list(store, "ny/2022/60cm/rgb/40073/")
+out = result.collect()
+len(out)
 
 
 def test_stream_sync():