pyo3-geoarrow 0.7.0

//! Python wrapper for GeoArrow scalar geometries.

#[cfg(feature = "geozero")]
mod bounding_rect;

use std::io::Write;
use std::sync::Arc;

use geoarrow_array::cast::AsGeoArrowArray;
use geoarrow_array::{GeoArrowArray, GeoArrowArrayAccessor, downcast_geoarrow_array};
use geoarrow_expr_geo::util::to_geo::geometry_to_geo;
use geoarrow_schema::GeoArrowType;
use geoarrow_schema::error::GeoArrowError;
use pyo3::exceptions::{PyIOError, PyValueError};
use pyo3::intern;
use pyo3::prelude::*;
use pyo3::types::{PyCapsule, PyTuple};
use pyo3_arrow::ffi::to_array_pycapsules;

use crate::PyGeoArray;
use crate::data_type::PyGeoType;
use crate::error::PyGeoArrowResult;
use crate::utils::text_repr::text_repr;

/// Python wrapper for a GeoArrow scalar geometry.
///
/// A scalar geometry represents a single geometry value. Internally, it is modeled as a
/// GeoArrow array of length 1, allowing it to use the same zero-copy Arrow C Data Interface
/// for interoperability.
#[pyclass(module = "geoarrow.rust.core", name = "GeoScalar", subclass, frozen)]
pub struct PyGeoScalar(Arc<dyn GeoArrowArray>);

impl PyGeoScalar {
    /// Create a new [`PyGeoScalar`] from a GeoArrow array.
    ///
    /// The array must have length 1.
    ///
    /// # Errors
    ///
    /// Returns a `PyValueError` if the array length is not 1.
    pub fn try_new(array: Arc<dyn GeoArrowArray>) -> PyGeoArrowResult<Self> {
        if array.len() != 1 {
            Err(
                PyValueError::new_err("Scalar geometry must be backed by an array of length 1.")
                    .into(),
            )
        } else {
            Ok(Self(array))
        }
    }

    /// Access a reference to the underlying GeoArrow array.
    pub fn inner(&self) -> &Arc<dyn GeoArrowArray> {
        &self.0
    }

    /// Consume this wrapper and return the underlying GeoArrow array.
    pub fn into_inner(self) -> Arc<dyn GeoArrowArray> {
        self.0
    }
}

#[pymethods]
impl PyGeoScalar {
    #[pyo3(signature = (requested_schema=None))]
    fn __arrow_c_array__<'py>(
        &'py self,
        py: Python<'py>,
        requested_schema: Option<Bound<'py, PyCapsule>>,
    ) -> PyGeoArrowResult<Bound<'py, PyTuple>> {
        let field = Arc::new(self.0.data_type().to_field("", true));
        let array = self.0.to_array_ref();
        Ok(to_array_pycapsules(py, field, &array, requested_schema)?)
    }

    fn __eq__(&self, other: &Bound<PyAny>) -> bool {
        // Do extraction within body because `__eq__` should never raise an exception.
        if let Ok(other) = other.extract::<Self>() {
            self.0.data_type() == other.0.data_type()
                && self.0.to_array_ref() == other.0.to_array_ref()
        } else {
            false
        }
    }

    #[cfg(feature = "geojson")]
    #[getter]
    fn __geo_interface__<'py>(&'py self, py: Python<'py>) -> PyGeoArrowResult<Bound<'py, PyAny>> {
        let geojson_geometry = scalar_to_geojson(&self.0)?;
        let geojson_string = serde_json::to_string(&geojson_geometry)?;
        let json_mod = py.import(intern!(py, "json"))?;
        Ok(json_mod.call_method1(intern!(py, "loads"), (geojson_string,))?)
    }

    #[cfg(feature = "geozero")]
    fn _repr_svg_(&self) -> PyGeoArrowResult<String> {
        use geozero::FeatureProcessor;

        use crate::scalar::bounding_rect::bounding_rect;

        let bounds = bounding_rect(&self.0)?.unwrap_or_default();
        let mut min_x = bounds.minx();
        let mut min_y = bounds.miny();
        let mut max_x = bounds.maxx();
        let mut max_y = bounds.maxy();

        let mut svg_data = Vec::new();
        // Passing `true` to `invert_y` is necessary to match Shapely's _repr_svg_
        let mut svg = geozero::svg::SvgWriter::new(&mut svg_data, true);

        // Expand box by 10% for readability
        min_x -= (max_x - min_x) * 0.05;
        min_y -= (max_y - min_y) * 0.05;
        max_x += (max_x - min_x) * 0.05;
        max_y += (max_y - min_y) * 0.05;

        svg.set_dimensions(min_x, min_y, max_x, max_y, 300, 300);

        // This sequence is necessary so that the SvgWriter writes the header. See
        // https://github.com/georust/geozero/blob/6c820ad7a0cac8c864058c783f548407427712d3/geozero/src/svg/mod.rs#L51-L58
        svg.dataset_begin(None)
            .map_err(|err| GeoArrowError::External(Box::new(err)))?;
        svg.feature_begin(0)
            .map_err(|err| GeoArrowError::External(Box::new(err)))?;
        if self.0.is_valid(0) {
            process_svg_geom(&self.0, &mut svg)
                .map_err(|err| GeoArrowError::External(Box::new(err)))?;
        }
        svg.feature_end(0)
            .map_err(|err| GeoArrowError::External(Box::new(err)))?;
        svg.dataset_end()
            .map_err(|err| GeoArrowError::External(Box::new(err)))?;

        let string =
            String::from_utf8(svg_data).map_err(|err| PyIOError::new_err(err.to_string()))?;
        Ok(string)
    }

    fn __repr__(&self) -> String {
        // TODO: print WKT representation
        format!("GeoScalar({})", text_repr(&self.0.data_type()))
    }

    #[getter]
    fn is_null(&self) -> bool {
        self.0.is_null(0)
    }

    #[getter]
    fn r#type(&self) -> PyGeoType {
        self.0.data_type().into()
    }
}

impl<'py> FromPyObject<'_, 'py> for PyGeoScalar {
    type Error = PyErr;

    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {
        Ok(Self::try_new(ob.extract::<PyGeoArray>()?.into_inner())?)
    }
}

#[cfg(feature = "geozero")]
fn process_svg_geom<W: Write>(
    arr: &dyn GeoArrowArray,
    svg: &mut geozero::svg::SvgWriter<W>,
) -> geozero::error::Result<()> {
    use GeoArrowType::*;
    use geozero::GeozeroGeometry;
    match arr.data_type() {
        Point(_) => arr.as_point().process_geom(svg),
        LineString(_) => arr.as_line_string().process_geom(svg),
        Polygon(_) => arr.as_polygon().process_geom(svg),
        MultiPoint(_) => arr.as_multi_point().process_geom(svg),
        MultiLineString(_) => arr.as_multi_line_string().process_geom(svg),
        MultiPolygon(_) => arr.as_multi_polygon().process_geom(svg),
        GeometryCollection(_) => arr.as_geometry_collection().process_geom(svg),
        Geometry(_) => arr.as_geometry().process_geom(svg),
        Rect(_) => arr.as_rect().process_geom(svg),
        Wkb(_) => arr.as_wkb::<i32>().process_geom(svg),
        LargeWkb(_) => arr.as_wkb::<i64>().process_geom(svg),
        WkbView(_) => arr.as_wkb_view().process_geom(svg),
        Wkt(_) => arr.as_wkt::<i32>().process_geom(svg),
        LargeWkt(_) => arr.as_wkt::<i64>().process_geom(svg),
        WktView(_) => arr.as_wkt_view().process_geom(svg),
    }
}

#[cfg(feature = "geojson")]
fn scalar_to_geojson(scalar: &dyn GeoArrowArray) -> PyGeoArrowResult<geojson::Geometry> {
    downcast_geoarrow_array!(scalar, impl_to_geojson)
}

#[cfg(feature = "geojson")]
fn impl_to_geojson<'a>(
    array: &'a impl GeoArrowArrayAccessor<'a>,
) -> PyGeoArrowResult<geojson::Geometry> {
    let geo_geom = geometry_to_geo(&array.value(0)?)?;
    Ok((&geo_geom).into())
}