feat: add walk

Author: gadomski

Cargo.lock

@@ -39,6 +39,7 @@ stac-cli = { git = "https://github.com/stac-utils/stac-rs", features = [
 stac-duckdb = { git = "https://github.com/stac-utils/stac-rs" }
 thiserror = "2.0.11"
 tokio = { version = "1.43.0", features = ["rt-multi-thread"] }
+pyo3-log = "0.12.1"
 
 [patch.crates-io]
 duckdb = { git = "https://github.com/duckdb/duckdb-rs", rev = "5eeb1f01c278790ce1e2d24045f0096e9e2528e4" }

Cargo.toml

@@ -18,6 +18,8 @@ type Result<T> = std::result::Result<T, Error>;
 
 #[pymodule]
 fn stacrs(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+    pyo3_log::init();
+
     m.add("StacrsError", py.get_type::<error::StacrsError>())?;
 
     m.add_class::<duckdb::DuckdbClient>()?;

src/lib.rs

@@ -1,25 +1,31 @@
 use crate::{Error, Json};
 use pyo3::{pyfunction, types::PyAny, Bound, PyResult, Python};
-use stac::{Format, Value};
+use stac::{Format, Link, Links, SelfHref, Value};
 
 #[pyfunction]
-#[pyo3(signature = (href, *, format=None, options=None))]
+#[pyo3(signature = (href, *, format=None, options=None, set_self_link=true))]
 pub fn read(
     py: Python<'_>,
     href: String,
     format: Option<String>,
     options: Option<Vec<(String, String)>>,
+    set_self_link: bool,
 ) -> PyResult<Bound<'_, PyAny>> {
     let format = format
         .and_then(|f| f.parse::<Format>().ok())
         .or_else(|| Format::infer_from_href(&href))
         .unwrap_or_default();
     let options = options.unwrap_or_default();
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
-        let value = format
+        let mut value = format
             .get_opts::<Value, _, _, _>(href, options)
             .await
             .map_err(Error::from)?;
+        if set_self_link {
+            if let Some(href) = value.self_href().cloned() {
+                value.set_link(Link::self_(href));
+            }
+        }
         Ok(Json(value))
     })
 }

src/read.rs

@@ -1,45 +1,55 @@
-use std::collections::VecDeque;
-
 use crate::{Error, Json, Result};
 use pyo3::{
     exceptions::PyStopAsyncIteration, pyclass, pyfunction, pymethods, types::PyDict, Bound, Py,
     PyAny, PyResult, Python,
 };
-use stac::{Container, Node, Value};
+use stac::{Container, Item, Links, Node, SelfHref, Value};
+use std::collections::VecDeque;
+use std::sync::Arc;
+use tokio::sync::Mutex;
 
 #[pyfunction]
 pub fn walk(container: Bound<'_, PyDict>) -> Result<Walk> {
-    let value: Value = pythonize::depythonize(&container)?;
+    let mut value: Value = pythonize::depythonize(&container)?;
+    if let Some(link) = value.link("self").cloned() {
+        *value.self_href_mut() = Some(link.href);
+    }
     let container: Container = value.try_into()?;
     let node = Node::from(container);
     let mut walks = VecDeque::new();
     walks.push_back(node);
-    Ok(Walk(walks))
+    Ok(Walk(Arc::new(Mutex::new(walks))))
 }
 
 #[pyclass]
-pub struct Walk(VecDeque<Node>);
+pub struct Walk(Arc<Mutex<VecDeque<Node>>>);
 
 #[pymethods]
 impl Walk {
     fn __aiter__(slf: Py<Self>) -> Py<Self> {
         slf
     }
 
-    fn __anext__<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        if let Some(node) = self.0.pop_front() {
-            pyo3_async_runtimes::tokio::future_into_py(py, async move {
-                let node = node.resolve().await.map_err(Error::from)?;
-                let mut children = Vec::with_capacity(node.children.len());
-                for child in node.children {
-                    children.push(Json(child.value.clone()));
-                    self.0.push_back(child);
-                }
-                let items: Vec<_> = node.items.iter().map(|item| Json(item.clone())).collect();
-                Ok((Json(node.value), children, items))
-            })
-        } else {
-            Err(PyStopAsyncIteration::new_err("walk complete"))
+    fn __anext__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let nodes = self.0.clone();
+        pyo3_async_runtimes::tokio::future_into_py(py, async move { next_walk(nodes).await })
+    }
+}
+
+type WalkStep = (Value, Vec<Container>, VecDeque<Item>);
+
+async fn next_walk(nodes: Arc<Mutex<VecDeque<Node>>>) -> PyResult<Json<WalkStep>> {
+    let mut nodes = nodes.lock().await;
+    if let Some(node) = nodes.pop_front() {
+        let mut node = node.resolve().await.map_err(Error::from)?;
+        let items = std::mem::take(&mut node.items);
+        let mut children = Vec::with_capacity(node.children.len());
+        for child in node.children {
+            children.push(child.value.clone());
+            nodes.push_back(child);
         }
+        Ok(Json((node.value.into(), children, items)))
+    } else {
+        Err(PyStopAsyncIteration::new_err("done walking"))
     }
 }

src/walk.rs

@@ -198,16 +198,18 @@ async def read(
     *,
     format: str | None = None,
     options: list[tuple[str, str]] | None = None,
+    set_self_link: bool = True,
 ) -> dict[str, Any]:
     """
     Reads STAC from a href.
 
     Args:
-        href (str): The href to write to
-        format (str | None): The input format. If not provided, will be inferred
+        href: The href to write to
+        format: The input format. If not provided, will be inferred
             from the href's extension.
-        options (list[tuple[str, str]] | None): Options for configuring an
+        options: Options for configuring an
             object store, e.g. your AWS credentials.
+        set_self_link: If True, set the `self` link to the value of `href`.
 
     Returns:
         The STAC value

stacrs.pyi

@@ -11,5 +11,5 @@ async def test_walk(examples: Path) -> None:
         all_children.extend(children)
         all_items.extend(items)
 
-    assert len(all_children) == 2
-    assert len(all_items) == 3
+    assert len(all_children) == 3
+    assert len(all_items) == 2