Optionally create root dir in LocalStore (#177)

Author: kylebarron

obstore/python/obstore/store/__init__.pyi

@@ -27,7 +27,11 @@ class LocalStore:
     ```
     """
     def __init__(
-        self, prefix: str | Path | None = None, *, automatic_cleanup: bool = False
+        self,
+        prefix: str | Path | None = None,
+        *,
+        automatic_cleanup: bool = False,
+        mkdir: bool = False,
     ) -> None:
         """Create a new LocalStore.
 
@@ -36,6 +40,7 @@ class LocalStore:
 
         Keyword Args:
             automatic_cleanup: if `True`, enables automatic cleanup of empty directories when deleting files. Defaults to False.
+            mkdir: if `True` and `prefix` is not `None`, the directory at `prefix` will attempt to be created. Note that this root directory will not be cleaned up, even if `automatic_cleanup` is `True`.
         """
     def __repr__(self) -> str: ...
     @classmethod

pyo3-object_store/src/local.rs

@@ -1,3 +1,4 @@
+use std::fs::create_dir_all;
 use std::sync::Arc;
 
 use object_store::local::LocalFileSystem;
@@ -29,12 +30,16 @@ impl PyLocalStore {
 #[pymethods]
 impl PyLocalStore {
     #[new]
-    #[pyo3(signature = (prefix = None, *, automatic_cleanup=false))]
+    #[pyo3(signature = (prefix = None, *, automatic_cleanup=false, mkdir=false))]
     fn py_new(
         prefix: Option<std::path::PathBuf>,
         automatic_cleanup: bool,
+        mkdir: bool,
     ) -> PyObjectStoreResult<Self> {
         let fs = if let Some(prefix) = prefix {
+            if mkdir {
+                create_dir_all(&prefix)?;
+            }
             LocalFileSystem::new_with_prefix(prefix)?
         } else {
             LocalFileSystem::new()

tests/store/test_local.py

@@ -1,3 +1,4 @@
+import tempfile
 from pathlib import Path
 
 import pytest
@@ -42,3 +43,14 @@ def test_local_from_url():
     url = f"file://{HERE.absolute()}//"
     with pytest.raises(GenericError):
         store = LocalStore.from_url(url)
+
+
+def test_create_prefix():
+    tmpdir = Path(tempfile.gettempdir()) / "abc"
+    assert not tmpdir.exists()
+    LocalStore(tmpdir, mkdir=True)
+    assert tmpdir.exists()
+
+    # Assert that mkdir=True works even when the dir already exists
+    LocalStore(tmpdir, mkdir=True)
+    assert tmpdir.exists()