feat: move in the search

Author: gadomski

Cargo.lock

@@ -30,7 +30,6 @@ stac = { features = [
 ], git = "https://github.com/stac-utils/stac-rs", branch = "main" }
 stac-api = { features = [
     "client",
-    "python",
 ], git = "https://github.com/stac-utils/stac-rs", branch = "main" }
 stac-cli = { git = "https://github.com/stac-utils/stac-rs", features = [
     "pgstac",

Cargo.toml

@@ -1,12 +1,14 @@
-use crate::Result;
+use crate::{
+    search::{PySortby, StringOrDict, StringOrList},
+    Result,
+};
 use pyo3::{
     exceptions::PyException,
     prelude::*,
     types::{PyDict, PyList},
     IntoPyObjectExt,
 };
 use pyo3_arrow::PyTable;
-use stac_api::python::{StringOrDict, StringOrList};
 use stac_duckdb::{Client, Config};
 use std::sync::Mutex;
 
@@ -16,11 +18,24 @@ pub struct DuckdbClient(Mutex<Client>);
 #[pymethods]
 impl DuckdbClient {
     #[new]
-    #[pyo3(signature = (use_s3_credential_chain=true, use_hive_partitioning=false))]
-    fn new(use_s3_credential_chain: bool, use_hive_partitioning: bool) -> Result<DuckdbClient> {
+    #[pyo3(signature = (*, use_s3_credential_chain=true, use_azure_credential_chain=true, use_httpfs=true, use_hive_partitioning=false, install_extensions=true, custom_extension_repository=None, extension_directory=None))]
+    fn new(
+        use_s3_credential_chain: bool,
+        use_azure_credential_chain: bool,
+        use_httpfs: bool,
+        use_hive_partitioning: bool,
+        install_extensions: bool,
+        custom_extension_repository: Option<String>,
+        extension_directory: Option<String>,
+    ) -> Result<DuckdbClient> {
         let config = Config {
             use_s3_credential_chain,
+            use_azure_credential_chain,
+            use_httpfs,
             use_hive_partitioning,
+            install_extensions,
+            custom_extension_repository,
+            extension_directory,
             convert_wkb: true,
         };
         let client = Client::with_config(config)?;
@@ -40,12 +55,12 @@ impl DuckdbClient {
         datetime: Option<String>,
         include: Option<StringOrList>,
         exclude: Option<StringOrList>,
-        sortby: Option<StringOrList>,
+        sortby: Option<PySortby<'py>>,
         filter: Option<StringOrDict>,
         query: Option<Bound<'py, PyDict>>,
         kwargs: Option<Bound<'py, PyDict>>,
     ) -> Result<Bound<'py, PyDict>> {
-        let search = stac_api::python::search(
+        let search = crate::search::build(
             intersects,
             ids,
             collections,
@@ -84,12 +99,12 @@ impl DuckdbClient {
         datetime: Option<String>,
         include: Option<StringOrList>,
         exclude: Option<StringOrList>,
-        sortby: Option<StringOrList>,
+        sortby: Option<PySortby<'py>>,
         filter: Option<StringOrDict>,
         query: Option<Bound<'py, PyDict>>,
         kwargs: Option<Bound<'py, PyDict>>,
     ) -> Result<PyObject> {
-        let search = stac_api::python::search(
+        let search = crate::search::build(
             intersects,
             ids,
             collections,

src/duckdb.rs

@@ -1,10 +1,10 @@
 use crate::{Error, Json, Result};
-use pyo3::{prelude::*, types::PyDict};
+use geojson::Geometry;
+use pyo3::prelude::*;
+use pyo3::{exceptions::PyValueError, types::PyDict, Bound, FromPyObject, PyErr, PyResult};
+use stac::Bbox;
 use stac::Format;
-use stac_api::{
-    python::{StringOrDict, StringOrList},
-    Search,
-};
+use stac_api::{Fields, Filter, Items, Search, Sortby};
 
 #[pyfunction]
 #[pyo3(signature = (href, *, intersects=None, ids=None, collections=None, max_items=None, limit=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None, use_duckdb=None, **kwargs))]
@@ -21,13 +21,13 @@ pub fn search<'py>(
     datetime: Option<String>,
     include: Option<StringOrList>,
     exclude: Option<StringOrList>,
-    sortby: Option<StringOrList>,
+    sortby: Option<PySortby<'py>>,
     filter: Option<StringOrDict>,
     query: Option<Bound<'py, PyDict>>,
     use_duckdb: Option<bool>,
     kwargs: Option<Bound<'_, PyDict>>,
 ) -> PyResult<Bound<'py, PyAny>> {
-    let search = stac_api::python::search(
+    let search = build(
         intersects,
         ids,
         collections,
@@ -72,15 +72,15 @@ pub fn search_to<'py>(
     datetime: Option<String>,
     include: Option<StringOrList>,
     exclude: Option<StringOrList>,
-    sortby: Option<StringOrList>,
+    sortby: Option<PySortby<'py>>,
     filter: Option<StringOrDict>,
     query: Option<Bound<'py, PyDict>>,
     format: Option<String>,
     options: Option<Vec<(String, String)>>,
     use_duckdb: Option<bool>,
     kwargs: Option<Bound<'_, PyDict>>,
 ) -> PyResult<Bound<'py, PyAny>> {
-    let search = stac_api::python::search(
+    let search = build(
         intersects,
         ids,
         collections,
@@ -150,3 +150,144 @@ async fn search_api(
     let value = stac_api::client::search(&href, search, max_items).await?;
     Ok(value)
 }
+
+/// Creates a [Search] from Python arguments.
+#[allow(clippy::too_many_arguments)]
+pub fn build<'py>(
+    intersects: Option<StringOrDict<'py>>,
+    ids: Option<StringOrList>,
+    collections: Option<StringOrList>,
+    limit: Option<u64>,
+    bbox: Option<Vec<f64>>,
+    datetime: Option<String>,
+    include: Option<StringOrList>,
+    exclude: Option<StringOrList>,
+    sortby: Option<PySortby<'py>>,
+    filter: Option<StringOrDict<'py>>,
+    query: Option<Bound<'py, PyDict>>,
+    kwargs: Option<Bound<'py, PyDict>>,
+) -> PyResult<Search> {
+    let mut fields = Fields::default();
+    if let Some(include) = include {
+        fields.include = include.into();
+    }
+    if let Some(exclude) = exclude {
+        fields.exclude = exclude.into();
+    }
+    let fields = if fields.include.is_empty() && fields.exclude.is_empty() {
+        None
+    } else {
+        Some(fields)
+    };
+    let query = query
+        .map(|query| pythonize::depythonize(&query))
+        .transpose()?;
+    let bbox = bbox.map(Bbox::try_from).transpose().map_err(Error::from)?;
+    let sortby: Vec<Sortby> = sortby
+        .map(|sortby| match sortby {
+            PySortby::ListOfDicts(list) => list
+                .into_iter()
+                .map(|d| pythonize::depythonize(&d).map_err(Error::from))
+                .collect::<Result<Vec<_>>>(),
+            PySortby::ListOfStrings(list) => list
+                .into_iter()
+                .map(|s| Ok(s.parse().unwrap())) // infallible
+                .collect::<Result<Vec<_>>>(),
+            PySortby::String(s) => Ok(vec![s.parse().unwrap()]),
+        })
+        .transpose()?
+        .unwrap_or_default();
+    let filter = filter
+        .map(|filter| match filter {
+            StringOrDict::Dict(cql_json) => pythonize::depythonize(&cql_json).map(Filter::Cql2Json),
+            StringOrDict::String(cql2_text) => Ok(Filter::Cql2Text(cql2_text)),
+        })
+        .transpose()?;
+    let filter = filter
+        .map(|filter| filter.into_cql2_json())
+        .transpose()
+        .map_err(Error::from)?;
+    let mut items = Items {
+        limit,
+        bbox,
+        datetime,
+        query,
+        fields,
+        sortby,
+        filter,
+        ..Default::default()
+    };
+    if let Some(kwargs) = kwargs {
+        items.additional_fields = pythonize::depythonize(&kwargs)?;
+    }
+
+    let intersects = intersects
+        .map(|intersects| match intersects {
+            StringOrDict::Dict(json) => pythonize::depythonize(&json)
+                .map_err(PyErr::from)
+                .and_then(|json| {
+                    Geometry::from_json_object(json)
+                        .map_err(|err| PyValueError::new_err(err.to_string()))
+                }),
+            StringOrDict::String(s) => s
+                .parse::<Geometry>()
+                .map_err(|err| PyValueError::new_err(err.to_string())),
+        })
+        .transpose()?;
+    let ids = ids.map(|ids| ids.into()).unwrap_or_default();
+    let collections = collections.map(|ids| ids.into()).unwrap_or_default();
+    Ok(Search {
+        items,
+        intersects,
+        ids,
+        collections,
+    })
+}
+
+/// A string or dictionary.
+///
+/// Used for the CQL2 filter argument and for intersects.
+#[derive(Debug, FromPyObject)]
+pub enum StringOrDict<'py> {
+    /// Text
+    String(String),
+
+    /// Json
+    Dict(Bound<'py, PyDict>),
+}
+
+/// A string or a list.
+///
+/// Used for collections, ids, etc.
+#[derive(Debug, FromPyObject)]
+pub enum StringOrList {
+    /// A string.
+    String(String),
+
+    /// A list.
+    List(Vec<String>),
+}
+
+/// A sortby structure.
+///
+/// This can be a string, a list of strings, or a list of dictionaries.
+#[derive(Debug, FromPyObject)]
+pub enum PySortby<'py> {
+    /// A string.
+    String(String),
+
+    /// A list.
+    ListOfStrings(Vec<String>),
+
+    /// A list.
+    ListOfDicts(Vec<Bound<'py, PyDict>>),
+}
+
+impl From<StringOrList> for Vec<String> {
+    fn from(value: StringOrList) -> Vec<String> {
+        match value {
+            StringOrList::List(list) => list,
+            StringOrList::String(s) => vec![s],
+        }
+    }
+}