Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),

## [Unreleased]

### Added

- Construct `stac_api::Search` (moved from `stac_api` crate) ([#81](https://github.com/stac-utils/stacrs/pull/81))

### Fixed

- Swallow broken pipe errors ([#73](https://github.com/stac-utils/stacrs/pull/73))
Expand Down
20 changes: 9 additions & 11 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 0 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
31 changes: 23 additions & 8 deletions src/duckdb.rs
Original file line number Diff line number Diff line change
@@ -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;

Expand All @@ -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)?;
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
159 changes: 150 additions & 9 deletions src/search.rs
Original file line number Diff line number Diff line change
@@ -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))]
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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],
}
}
}
Loading
Loading