From 74c5efcc7b4b8fd03efeaf4bb6098f11e31b8306 Mon Sep 17 00:00:00 2001 From: Michael Kirk Date: Thu, 1 Sep 2022 10:47:13 -0700 Subject: [PATCH] Add FeatureWriter to stream writes --- CHANGES.md | 3 + examples/stream_reader_writer.rs | 45 ++++ src/de.rs | 5 +- src/errors.rs | 2 + src/feature_writer.rs | 441 +++++++++++++++++++++++++++++++ src/lib.rs | 4 +- 6 files changed, 497 insertions(+), 3 deletions(-) create mode 100644 examples/stream_reader_writer.rs create mode 100644 src/feature_writer.rs diff --git a/CHANGES.md b/CHANGES.md index 95f0e6e..1ddab60 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -26,6 +26,9 @@ geojson::ser::to_feature_collection_string(&my_structs).unwrap(); ``` * PR: +* Added `geojson::{FeatureReader, FeatureWriter}` to stream the reading/writing of your custom struct to and from GeoJSON. + * PR: + * PR: * Added IntoIter implementation for FeatureCollection. * * Added `geojson::Result`. diff --git a/examples/stream_reader_writer.rs b/examples/stream_reader_writer.rs new file mode 100644 index 0000000..4ae17fd --- /dev/null +++ b/examples/stream_reader_writer.rs @@ -0,0 +1,45 @@ +use geojson::{de::deserialize_geometry, ser::serialize_geometry, FeatureReader, FeatureWriter}; + +use serde::{Deserialize, Serialize}; +use std::error::Error; +use std::fs::File; +use std::io::{BufReader, BufWriter}; + +#[cfg(not(feature = "geo-types"))] +fn main() -> Result<(), Box> { + panic!("this example requires geo-types") +} + +#[cfg(feature = "geo-types")] +fn main() -> Result<(), Box> { + #[derive(Serialize, Deserialize)] + struct Country { + #[serde( + serialize_with = "serialize_geometry", + deserialize_with = "deserialize_geometry" + )] + geometry: geo_types::Geometry, + name: String, + } + + let reader = { + let file_reader = BufReader::new(File::open("tests/fixtures/countries.geojson")?); + FeatureReader::from_reader(file_reader) + }; + + let mut writer = { + let file_writer = BufWriter::new(File::create("example-output-contries.geojson")?); + FeatureWriter::from_writer(file_writer) + }; + + let mut country_count = 0; + for country in reader.deserialize::()? { + let country = country?; + country_count += 1; + + writer.serialize(&country)?; + } + + assert_eq!(country_count, 180); + Ok(()) +} diff --git a/src/de.rs b/src/de.rs index 0489ddf..3023732 100644 --- a/src/de.rs +++ b/src/de.rs @@ -429,7 +429,6 @@ pub(crate) mod tests { use crate::JsonValue; - use serde::Deserialize; use serde_json::json; pub(crate) fn feature_collection() -> JsonValue { @@ -491,6 +490,8 @@ pub(crate) mod tests { mod geo_types_tests { use super::*; + use serde::Deserialize; + #[test] fn geometry_field() { #[derive(Deserialize)] @@ -592,7 +593,7 @@ pub(crate) mod tests { #[test] fn roundtrip() { use crate::ser::serialize_geometry; - use serde::Serialize; + use serde::{Deserialize, Serialize}; #[derive(Serialize, Deserialize)] struct MyStruct { diff --git a/src/errors.rs b/src/errors.rs index ed1f5b4..ace2b2e 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -15,6 +15,8 @@ pub enum Error { /// This was previously `GeoJsonUnknownType`, but has been split for clarity #[error("Expected a Feature, FeatureCollection, or Geometry, but got an empty type")] EmptyType, + #[error("invalid writer state: {0}")] + InvalidWriterState(&'static str), #[error("IO Error: {0}")] Io(std::io::Error), /// This was previously `GeoJsonUnknownType`, but has been split for clarity diff --git a/src/feature_writer.rs b/src/feature_writer.rs new file mode 100644 index 0000000..527ef31 --- /dev/null +++ b/src/feature_writer.rs @@ -0,0 +1,441 @@ +use crate::ser::to_feature_writer; +use crate::{Error, Feature, Result}; + +use serde::Serialize; +use std::io::Write; + +#[derive(PartialEq)] +enum State { + New, + Started, + Finished, +} + +/// Write Features to a FeatureCollection +pub struct FeatureWriter { + writer: W, + state: State, +} + +impl FeatureWriter { + /// Create a FeatureWriter from the given `writer`. + /// + /// To append features from your custom structs, use [`FeatureWriter::serialize`]. + /// + /// To append features from [`geojson::Feature`] use [`FeatureWriter::write_feature`]. + pub fn from_writer(writer: W) -> Self { + Self { + writer, + state: State::New, + } + } + + /// Write a [`crate::Feature`] struct to the output stream. If you'd like to + /// serialize your own custom structs, see [`FeatureWriter::serialize`] instead. + pub fn write_feature(&mut self, feature: &Feature) -> Result<()> { + match self.state { + State::Finished => { + return Err(Error::InvalidWriterState( + "cannot write another Feature when writer has already finished", + )) + } + State::New => { + self.write_prefix()?; + self.state = State::Started; + } + State::Started => { + self.write_str(",")?; + } + } + serde_json::to_writer(&mut self.writer, feature)?; + Ok(()) + } + + /// Serialize your own custom struct to the features of a FeatureCollection using the + /// [`serde`] crate. + /// + /// # Examples + /// + /// Your struct must implement or derive [`serde::Serialize`]. + /// + /// If you have enabled the `geo-types` feature, which is enabled by default, you can + /// serialize directly from a useful geometry type. + /// + /// ```rust,ignore + /// use geojson::{FeatureWriter, ser::serialize_geometry}; + /// + /// #[derive(serde::Serialize)] + /// struct MyStruct { + /// #[serde(serialize_with = "serialize_geometry")] + /// geometry: geo_types::Point, + /// name: String, + /// age: u64, + /// } + /// ``` + /// + /// Then you can serialize the FeatureCollection directly from your type. + #[cfg_attr(feature = "geo-types", doc = "```")] + #[cfg_attr(not(feature = "geo-types"), doc = "```ignore")] + /// # + /// # use geojson::{FeatureWriter, ser::serialize_geometry}; + /// # + /// # #[derive(serde::Serialize)] + /// # struct MyStruct { + /// # #[serde(serialize_with = "serialize_geometry")] + /// # geometry: geo_types::Point, + /// # name: String, + /// # age: u64, + /// # } + /// + /// let dinagat = MyStruct { + /// geometry: geo_types::point!(x: 125.6, y: 10.1), + /// name: "Dinagat Islands".to_string(), + /// age: 123 + /// }; + /// + /// let neverland = MyStruct { + /// geometry: geo_types::point!(x: 2.3, y: 4.5), + /// name: "Neverland".to_string(), + /// age: 456 + /// }; + /// + /// let mut output: Vec = vec![]; + /// { + /// let io_writer = std::io::BufWriter::new(&mut output); + /// let mut feature_writer = FeatureWriter::from_writer(io_writer); + /// feature_writer.serialize(&dinagat).unwrap(); + /// feature_writer.serialize(&neverland).unwrap(); + /// } + /// + /// let expected_output = r#"{ + /// "type": "FeatureCollection", + /// "features": [ + /// { + /// "type": "Feature", + /// "geometry": { "type": "Point", "coordinates": [125.6, 10.1] }, + /// "properties": { + /// "name": "Dinagat Islands", + /// "age": 123 + /// } + /// }, + /// { + /// "type": "Feature", + /// "geometry": { "type": "Point", "coordinates": [2.3, 4.5] }, + /// "properties": { + /// "name": "Neverland", + /// "age": 456 + /// } + /// } + /// ] + /// }"#.as_bytes(); + /// + /// fn assert_eq_json(bytes_1: &[u8], bytes_2: &[u8]) { + /// // check for semantic equality, discarding any formatting/whitespace differences + /// let json_1: serde_json::Value = serde_json::from_slice(bytes_1).unwrap(); + /// let json_2: serde_json::Value = serde_json::from_slice(bytes_2).unwrap(); + /// assert_eq!(json_1, json_2); + /// } + /// + /// assert_eq_json(expected_output, &output); + /// ``` + /// + /// If you're not using [`geo-types`](geo_types), you can deserialize to a `geojson::Geometry` instead. + /// ```rust,ignore + /// use serde::Deserialize; + /// #[derive(Deserialize)] + /// struct MyStruct { + /// geometry: geojson::Geometry, + /// name: String, + /// age: u64, + /// } + /// ``` + pub fn serialize(&mut self, value: &S) -> Result<()> { + match self.state { + State::Finished => { + return Err(Error::InvalidWriterState( + "cannot serialize another record when writer has already finished", + )) + } + State::New => { + self.write_prefix()?; + self.state = State::Started; + } + State::Started => { + self.write_str(",")?; + } + } + to_feature_writer(&mut self.writer, value) + } + + /// Writes the closing syntax for the FeatureCollection. + /// + /// You shouldn't normally need to call this manually, as the writer will close itself upon + /// being dropped. + pub fn finish(&mut self) -> Result<()> { + match self.state { + State::Finished => { + return Err(Error::InvalidWriterState( + "cannot finish writer - it's already finished", + )) + } + State::New => { + self.state = State::Finished; + self.write_prefix()?; + self.write_suffix()?; + } + State::Started => { + self.state = State::Finished; + self.write_suffix()?; + } + } + Ok(()) + } + + /// Flush the underlying writer buffer. + /// + /// You shouldn't normally need to call this manually, as the writer will flush itself upon + /// being dropped. + pub fn flush(&mut self) -> Result<()> { + Ok(self.writer.flush()?) + } + + fn write_prefix(&mut self) -> Result<()> { + self.write_str(r#"{ "type": "FeatureCollection", "features": ["#) + } + + fn write_suffix(&mut self) -> Result<()> { + self.write_str("]}") + } + + fn write_str(&mut self, text: &str) -> Result<()> { + self.writer.write_all(text.as_bytes())?; + Ok(()) + } +} + +impl Drop for FeatureWriter { + fn drop(&mut self) { + if self.state != State::Finished { + _ = self.finish().map_err(|e| { + log::error!("FeatureWriter errored while finishing in Drop impl. To handle errors like this, explicitly call `FeatureWriter::finish`. Error: {}", e); + }); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::JsonValue; + + use serde_json::json; + + // an example struct that we want to serialize + #[derive(Serialize)] + struct MyRecord { + geometry: crate::Geometry, + name: String, + age: u64, + } + + #[test] + fn write_empty() { + let mut buffer: Vec = vec![]; + { + let mut writer = FeatureWriter::from_writer(&mut buffer); + writer.finish().unwrap(); + } + + let expected = json!({ + "type": "FeatureCollection", + "features": [] + }); + + let actual_json: JsonValue = serde_json::from_slice(&buffer).unwrap(); + assert_eq!(actual_json, expected); + } + + #[test] + fn finish_on_drop() { + let mut buffer: Vec = vec![]; + { + _ = FeatureWriter::from_writer(&mut buffer); + } + + let expected = json!({ + "type": "FeatureCollection", + "features": [] + }); + + let actual_json: JsonValue = serde_json::from_slice(&buffer).unwrap(); + assert_eq!(actual_json, expected); + } + + #[test] + fn write_feature() { + let mut buffer: Vec = vec![]; + { + let mut writer = FeatureWriter::from_writer(&mut buffer); + + let record_1 = { + let mut props = serde_json::Map::new(); + props.insert("name".to_string(), "Mishka".into()); + props.insert("age".to_string(), 12.into()); + + Feature { + bbox: None, + geometry: Some(crate::Geometry::from(crate::Value::Point(vec![1.1, 1.2]))), + id: None, + properties: Some(props), + foreign_members: None, + } + }; + + let record_2 = { + let mut props = serde_json::Map::new(); + props.insert("name".to_string(), "Jane".into()); + props.insert("age".to_string(), 22.into()); + + Feature { + bbox: None, + geometry: Some(crate::Geometry::from(crate::Value::Point(vec![2.1, 2.2]))), + id: None, + properties: Some(props), + foreign_members: None, + } + }; + + writer.write_feature(&record_1).unwrap(); + writer.write_feature(&record_2).unwrap(); + writer.flush().unwrap(); + } + + let expected = json!({ + "type": "FeatureCollection", + "features": [ + { + "type": "Feature", + "geometry": { "type": "Point", "coordinates": [1.1, 1.2] }, + "properties": { "name": "Mishka", "age": 12 + } + }, + { + "type": "Feature", + "geometry": { "type": "Point", "coordinates": [2.1, 2.2] }, + "properties": { + "name": "Jane", + "age": 22 + } + } + ] + }); + + let actual_json: JsonValue = serde_json::from_slice(&buffer).expect("valid json"); + assert_eq!(actual_json, expected) + } + + #[test] + fn serialize() { + let mut buffer: Vec = vec![]; + { + let mut writer = FeatureWriter::from_writer(&mut buffer); + let record_1 = MyRecord { + geometry: crate::Geometry::from(crate::Value::Point(vec![1.1, 1.2])), + name: "Mishka".to_string(), + age: 12, + }; + let record_2 = MyRecord { + geometry: crate::Geometry::from(crate::Value::Point(vec![2.1, 2.2])), + name: "Jane".to_string(), + age: 22, + }; + writer.serialize(&record_1).unwrap(); + writer.serialize(&record_2).unwrap(); + writer.flush().unwrap(); + } + + let expected = json!({ + "type": "FeatureCollection", + "features": [ + { + "type": "Feature", + "geometry": { "type": "Point", "coordinates": [1.1, 1.2] }, + "properties": { "name": "Mishka", "age": 12 + } + }, + { + "type": "Feature", + "geometry": { "type": "Point", "coordinates": [2.1, 2.2] }, + "properties": { + "name": "Jane", + "age": 22 + } + } + ] + }); + + let actual_json: JsonValue = serde_json::from_slice(&buffer).expect("valid json"); + assert_eq!(actual_json, expected) + } + + #[cfg(feature = "geo-types")] + mod test_geo_types { + use super::*; + use crate::ser::serialize_geometry; + + // an example struct that we want to serialize + #[derive(Serialize)] + struct MyGeoRecord { + #[serde(serialize_with = "serialize_geometry")] + geometry: geo_types::Point, + name: String, + age: u64, + } + + #[test] + fn serialize_geo_types() { + let mut buffer: Vec = vec![]; + { + let mut writer = FeatureWriter::from_writer(&mut buffer); + let record_1 = MyGeoRecord { + geometry: geo_types::point!(x: 1.1, y: 1.2), + name: "Mishka".to_string(), + age: 12, + }; + let record_2 = MyGeoRecord { + geometry: geo_types::point!(x: 2.1, y: 2.2), + name: "Jane".to_string(), + age: 22, + }; + writer.serialize(&record_1).unwrap(); + writer.serialize(&record_2).unwrap(); + writer.flush().unwrap(); + } + + let expected = json!({ + "type": "FeatureCollection", + "features": [ + { + "type": "Feature", + "geometry": { "type": "Point", "coordinates": [1.1, 1.2] }, + "properties": { + "name": "Mishka", + "age": 12 + } + }, + { + "type": "Feature", + "geometry": { "type": "Point", "coordinates": [2.1, 2.2] }, + "properties": { + "name": "Jane", + "age": 22 + } + } + ] + }); + + let actual_json: JsonValue = serde_json::from_slice(&buffer).expect("valid json"); + assert_eq!(actual_json, expected) + } + } +} diff --git a/src/lib.rs b/src/lib.rs index d2e4bc9..3dee962 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -466,9 +466,11 @@ pub mod de; pub mod ser; mod feature_reader; - pub use feature_reader::FeatureReader; +mod feature_writer; +pub use feature_writer::FeatureWriter; + #[cfg(feature = "geo-types")] pub use conversion::quick_collection;