pyo3_object_store/
client.rs

1use std::collections::HashMap;
2use std::str::FromStr;
3
4use http::{HeaderMap, HeaderName, HeaderValue};
5use object_store::{ClientConfigKey, ClientOptions};
6use pyo3::exceptions::PyValueError;
7use pyo3::prelude::*;
8use pyo3::pybacked::{PyBackedBytes, PyBackedStr};
9use pyo3::types::{PyDict, PyString};
10
11use crate::config::PyConfigValue;
12use crate::error::PyObjectStoreError;
13
14/// A wrapper around `ClientConfigKey` that implements [`FromPyObject`].
15#[derive(Clone, Debug, PartialEq, Eq, Hash)]
16pub struct PyClientConfigKey(ClientConfigKey);
17
18impl<'py> FromPyObject<'_, 'py> for PyClientConfigKey {
19    type Error = PyErr;
20
21    fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult<Self> {
22        let s = obj.extract::<PyBackedStr>()?.to_lowercase();
23        let key = s.parse().map_err(PyObjectStoreError::ObjectStoreError)?;
24        Ok(Self(key))
25    }
26}
27
28impl<'py> IntoPyObject<'py> for PyClientConfigKey {
29    type Target = PyString;
30    type Output = Bound<'py, PyString>;
31    type Error = PyErr;
32
33    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
34        Ok(PyString::new(py, self.0.as_ref()))
35    }
36}
37
38impl<'py> IntoPyObject<'py> for &PyClientConfigKey {
39    type Target = PyString;
40    type Output = Bound<'py, PyString>;
41    type Error = PyErr;
42
43    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
44        Ok(PyString::new(py, self.0.as_ref()))
45    }
46}
47
48/// A wrapper around `ClientOptions` that implements [`FromPyObject`].
49#[derive(Clone, Debug, PartialEq)]
50pub struct PyClientOptions {
51    string_options: HashMap<PyClientConfigKey, PyConfigValue>,
52    default_headers: Option<PyHeaderMap>,
53}
54
55impl<'py> FromPyObject<'_, 'py> for PyClientOptions {
56    type Error = PyErr;
57
58    // Need custom extraction because default headers is not a string value
59    fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult<Self> {
60        let dict = obj.extract::<Bound<PyDict>>()?;
61        let mut string_options = HashMap::new();
62        let mut default_headers = None;
63
64        for (key, value) in dict.iter() {
65            if let Ok(key) = key.extract::<PyClientConfigKey>() {
66                string_options.insert(key, value.extract::<PyConfigValue>()?);
67            } else {
68                let key = key.extract::<PyBackedStr>()?;
69                if &key == "default_headers" {
70                    default_headers = Some(value.extract::<PyHeaderMap>()?);
71                } else {
72                    return Err(PyValueError::new_err(format!("Invalid key: {key}.")));
73                }
74            }
75        }
76
77        Ok(Self {
78            string_options,
79            default_headers,
80        })
81    }
82}
83
84impl<'py> IntoPyObject<'py> for PyClientOptions {
85    type Target = PyDict;
86    type Output = Bound<'py, PyDict>;
87    type Error = PyErr;
88
89    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
90        let dict = self.string_options.into_pyobject(py)?;
91        if let Some(headers) = self.default_headers {
92            dict.set_item("default_headers", headers)?;
93        }
94        Ok(dict)
95    }
96}
97
98impl<'py> IntoPyObject<'py> for &PyClientOptions {
99    type Target = PyDict;
100    type Output = Bound<'py, PyDict>;
101    type Error = PyErr;
102
103    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
104        let dict = (&self.string_options).into_pyobject(py)?;
105        if let Some(headers) = &self.default_headers {
106            dict.set_item("default_headers", headers)?;
107        }
108        Ok(dict.clone())
109    }
110}
111
112impl From<PyClientOptions> for ClientOptions {
113    fn from(value: PyClientOptions) -> Self {
114        let mut options = ClientOptions::new();
115        for (key, value) in value.string_options.into_iter() {
116            options = options.with_config(key.0, value.0);
117        }
118
119        if let Some(headers) = value.default_headers {
120            options = options.with_default_headers(headers.0);
121        }
122
123        options
124    }
125}
126
127#[derive(Clone, Debug, PartialEq)]
128struct PyHeaderMap(HeaderMap);
129
130impl<'py> FromPyObject<'_, 'py> for PyHeaderMap {
131    type Error = PyErr;
132
133    fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult<Self> {
134        let dict = obj.extract::<Bound<PyDict>>()?;
135        let mut header_map = HeaderMap::with_capacity(dict.len());
136        for (key, value) in dict.iter() {
137            let key = HeaderName::from_str(&key.extract::<PyBackedStr>()?)
138                .map_err(|err| PyValueError::new_err(err.to_string()))?;
139
140            // HTTP Header values can have non-ascii bytes, so first try to extract as bytes.
141            let value = if let Ok(value_bytes) = value.extract::<PyBackedBytes>() {
142                HeaderValue::from_bytes(&value_bytes)
143            } else {
144                HeaderValue::from_str(&value.extract::<PyBackedStr>()?)
145            }
146            .map_err(|err| PyValueError::new_err(err.to_string()))?;
147
148            header_map.insert(key, value);
149        }
150        Ok(Self(header_map))
151    }
152}
153
154impl<'py> IntoPyObject<'py> for PyHeaderMap {
155    type Target = PyDict;
156    type Output = Bound<'py, PyDict>;
157    type Error = PyErr;
158
159    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
160        let dict = PyDict::new(py);
161        for (key, value) in self.0.iter() {
162            dict.set_item(key.as_str(), value.as_bytes())?;
163        }
164        Ok(dict)
165    }
166}
167
168impl<'py> IntoPyObject<'py> for &PyHeaderMap {
169    type Target = PyDict;
170    type Output = Bound<'py, PyDict>;
171    type Error = PyErr;
172
173    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
174        let dict = PyDict::new(py);
175        for (key, value) in self.0.iter() {
176            dict.set_item(key.as_str(), value.as_bytes())?;
177        }
178        Ok(dict)
179    }
180}
pyo3_object_store/client.rs

pyo3_object_store/
client.rs
