pyo3_object_store/
local.rs

1use std::fs::create_dir_all;
2use std::sync::Arc;
3
4use object_store::local::LocalFileSystem;
5use object_store::ObjectStoreScheme;
6use pyo3::exceptions::PyValueError;
7use pyo3::prelude::*;
8use pyo3::types::{PyDict, PyTuple, PyType};
9use pyo3::{intern, IntoPyObjectExt};
10
11use crate::error::PyObjectStoreResult;
12use crate::PyUrl;
13
14#[derive(Clone, Debug, PartialEq)]
15struct LocalConfig {
16    prefix: Option<std::path::PathBuf>,
17    automatic_cleanup: bool,
18    mkdir: bool,
19}
20
21impl LocalConfig {
22    fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
23        let args = PyTuple::new(py, vec![self.prefix.clone()])?.into_bound_py_any(py)?;
24        let kwargs = PyDict::new(py);
25        kwargs.set_item(intern!(py, "automatic_cleanup"), self.automatic_cleanup)?;
26        kwargs.set_item(intern!(py, "mkdir"), self.mkdir)?;
27        PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?])
28    }
29}
30
31/// A Python-facing wrapper around a [`LocalFileSystem`].
32#[derive(Debug, Clone)]
33#[pyclass(name = "LocalStore", frozen, subclass)]
34pub struct PyLocalStore {
35    store: Arc<LocalFileSystem>,
36    config: LocalConfig,
37}
38
39impl AsRef<Arc<LocalFileSystem>> for PyLocalStore {
40    fn as_ref(&self) -> &Arc<LocalFileSystem> {
41        &self.store
42    }
43}
44
45impl PyLocalStore {
46    /// Consume self and return the underlying [`LocalFileSystem`].
47    pub fn into_inner(self) -> Arc<LocalFileSystem> {
48        self.store
49    }
50}
51
52#[pymethods]
53impl PyLocalStore {
54    #[new]
55    #[pyo3(signature = (prefix=None, *, automatic_cleanup=false, mkdir=false))]
56    fn new(
57        prefix: Option<std::path::PathBuf>,
58        automatic_cleanup: bool,
59        mkdir: bool,
60    ) -> PyObjectStoreResult<Self> {
61        let fs = if let Some(prefix) = &prefix {
62            if mkdir {
63                create_dir_all(prefix)?;
64            }
65            LocalFileSystem::new_with_prefix(prefix)?
66        } else {
67            LocalFileSystem::new()
68        };
69        let fs = fs.with_automatic_cleanup(automatic_cleanup);
70        Ok(Self {
71            store: Arc::new(fs),
72            config: LocalConfig {
73                prefix,
74                automatic_cleanup,
75                mkdir,
76            },
77        })
78    }
79
80    #[classmethod]
81    #[pyo3(signature = (url, *, automatic_cleanup=false, mkdir=false))]
82    pub(crate) fn from_url<'py>(
83        cls: &Bound<'py, PyType>,
84        url: PyUrl,
85        automatic_cleanup: bool,
86        mkdir: bool,
87    ) -> PyObjectStoreResult<Bound<'py, PyAny>> {
88        let url = url.into_inner();
89        let (scheme, path) = ObjectStoreScheme::parse(&url).map_err(object_store::Error::from)?;
90
91        if !matches!(scheme, ObjectStoreScheme::Local) {
92            return Err(PyValueError::new_err("Not a `file://` URL").into());
93        }
94
95        // The path returned by `ObjectStoreScheme::parse` strips the initial `/`, so we join it
96        // onto a root
97        // Hopefully this also works on Windows.
98        let root = std::path::Path::new("/");
99        let full_path = root.join(path.as_ref());
100
101        // Note: we pass **back** through Python so that if cls is a subclass, we instantiate the
102        // subclass
103        let kwargs = PyDict::new(cls.py());
104        kwargs.set_item("prefix", full_path)?;
105        kwargs.set_item("automatic_cleanup", automatic_cleanup)?;
106        kwargs.set_item("mkdir", mkdir)?;
107        Ok(cls.call((), Some(&kwargs))?)
108    }
109
110    fn __eq__(&self, other: &Bound<PyAny>) -> bool {
111        // Ensure we never error on __eq__ by returning false if the other object is not the same
112        // type
113        other
114            .cast::<PyLocalStore>()
115            .map(|other| self.config == other.get().config)
116            .unwrap_or(false)
117    }
118
119    fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
120        self.config.__getnewargs_ex__(py)
121    }
122
123    fn __repr__(&self) -> String {
124        if let Some(prefix) = &self.config.prefix {
125            format!("LocalStore(\"{}\")", prefix.display())
126        } else {
127            "LocalStore".to_string()
128        }
129    }
130
131    #[getter]
132    fn prefix<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
133        // Note: returning a std::path::Path or std::path::PathBuf converts back to a Python _str_
134        // not a Python _pathlib.Path_.
135        // So we manually convert to a pathlib.Path
136        if let Some(prefix) = &self.config.prefix {
137            let pathlib_mod = py.import(intern!(py, "pathlib"))?;
138            pathlib_mod.call_method1(intern!(py, "Path"), PyTuple::new(py, vec![prefix])?)
139        } else {
140            py.None().into_bound_py_any(py)
141        }
142    }
143}
pyo3_object_store/local.rs

pyo3_object_store/
local.rs
