Support parsing Z, M, and ZM

georust · Jul 17, 2024 · e1d838d · e1d838d
1 parent 2f1f44f
commit e1d838d
Show file tree

Hide file tree

Showing 11 changed files with 135 additions and 71 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -84,13 +84,13 @@ use std::str::FromStr;
 use num_traits::{Float, Num, NumCast};
 
 use crate::tokenizer::{PeekableTokens, Token, Tokens};
-use crate::types::GeometryCollection;
 use crate::types::LineString;
 use crate::types::MultiLineString;
 use crate::types::MultiPoint;
 use crate::types::MultiPolygon;
 use crate::types::Point;
 use crate::types::Polygon;
+use crate::types::{Dimension, GeometryCollection};
 
 mod to_wkt;
 mod tokenizer;
@@ -165,60 +165,31 @@ where
     ) -> Result<Self, &'static str> {
         match word {
             w if w.eq_ignore_ascii_case("POINT") => {
-                let x = <Point<T> as FromTokens<T>>::from_tokens_with_parens(tokens);
+                let x = <Point<T> as FromTokens<T>>::from_tokens_with_header(tokens);
                 x.map(|y| y.as_item())
             }
-            w if w.eq_ignore_ascii_case("POINTZ") => {
-                let x = <Point<T> as FromTokens<T>>::from_tokens_with_parens(tokens)?;
-                if let Some(coord) = &x.0 {
-                    if coord.z.is_none() {
-                        return Err("POINTZ must have a z-coordinate.");
-                    }
-                }
-                Ok(x.as_item())
-            }
-            w if w.eq_ignore_ascii_case("POINTM") => {
-                let mut x = <Point<T> as FromTokens<T>>::from_tokens_with_parens(tokens)?;
-                if let Some(coord) = &mut x.0 {
-                    if coord.z.is_none() {
-                        return Err("POINTM must have an m-coordinate.");
-                    } else {
-                        coord.m = coord.z.take();
-                    }
-                }
-                Ok(x.as_item())
-            }
-            w if w.eq_ignore_ascii_case("POINTZM") => {
-                let x = <Point<T> as FromTokens<T>>::from_tokens_with_parens(tokens)?;
-                if let Some(coord) = &x.0 {
-                    if coord.z.is_none() || coord.m.is_none() {
-                        return Err("POINTZM must have both a z- and m-coordinate");
-                    }
-                }
-                Ok(x.as_item())
-            }
             w if w.eq_ignore_ascii_case("LINESTRING") || w.eq_ignore_ascii_case("LINEARRING") => {
-                let x = <LineString<T> as FromTokens<T>>::from_tokens_with_parens(tokens);
+                let x = <LineString<T> as FromTokens<T>>::from_tokens_with_header(tokens);
                 x.map(|y| y.as_item())
             }
             w if w.eq_ignore_ascii_case("POLYGON") => {
-                let x = <Polygon<T> as FromTokens<T>>::from_tokens_with_parens(tokens);
+                let x = <Polygon<T> as FromTokens<T>>::from_tokens_with_header(tokens);
                 x.map(|y| y.as_item())
             }
             w if w.eq_ignore_ascii_case("MULTIPOINT") => {
-                let x = <MultiPoint<T> as FromTokens<T>>::from_tokens_with_parens(tokens);
+                let x = <MultiPoint<T> as FromTokens<T>>::from_tokens_with_header(tokens);
                 x.map(|y| y.as_item())
             }
             w if w.eq_ignore_ascii_case("MULTILINESTRING") => {
-                let x = <MultiLineString<T> as FromTokens<T>>::from_tokens_with_parens(tokens);
+                let x = <MultiLineString<T> as FromTokens<T>>::from_tokens_with_header(tokens);
                 x.map(|y| y.as_item())
             }
             w if w.eq_ignore_ascii_case("MULTIPOLYGON") => {
-                let x = <MultiPolygon<T> as FromTokens<T>>::from_tokens_with_parens(tokens);
+                let x = <MultiPolygon<T> as FromTokens<T>>::from_tokens_with_header(tokens);
                 x.map(|y| y.as_item())
             }
             w if w.eq_ignore_ascii_case("GEOMETRYCOLLECTION") => {
-                let x = <GeometryCollection<T> as FromTokens<T>>::from_tokens_with_parens(tokens);
+                let x = <GeometryCollection<T> as FromTokens<T>>::from_tokens_with_header(tokens);
                 x.map(|y| y.as_item())
             }
             _ => Err("Invalid type encountered"),
@@ -293,21 +264,64 @@ where
     }
 }
 
+fn infer_geom_dimension<T: WktNum + FromStr + Default>(
+    tokens: &mut PeekableTokens<T>,
+) -> Result<Dimension, &'static str> {
+    if let Some(Ok(c)) = tokens.peek() {
+        match c {
+            // If we match a word check if it's Z/M/ZM and consume the token from the stream
+            Token::Word(x) => match x.as_str() {
+                "Z" => {
+                    tokens.next().unwrap().unwrap();
+                    Ok(Dimension::XYZ)
+                }
+                "M" => {
+                    tokens.next().unwrap().unwrap();
+
+                    Ok(Dimension::XYM)
+                }
+                "ZM" => {
+                    tokens.next().unwrap().unwrap();
+                    Ok(Dimension::XYZM)
+                }
+                _ => Err("Unexpected word before open paren"),
+            },
+            // Not a word, e.g. an open paren
+            _ => Ok(Dimension::XY),
+        }
+    } else {
+        Err("End of stream")
+    }
+}
+
 trait FromTokens<T>: Sized + Default
 where
     T: WktNum + FromStr + Default,
 {
-    fn from_tokens(tokens: &mut PeekableTokens<T>) -> Result<Self, &'static str>;
+    fn from_tokens(tokens: &mut PeekableTokens<T>, dim: Dimension) -> Result<Self, &'static str>;
 
-    fn from_tokens_with_parens(tokens: &mut PeekableTokens<T>) -> Result<Self, &'static str> {
+    /// The preferred top-level FromTokens API, which additionally checks for the presence of Z, M,
+    /// and ZM in the token stream.
+    fn from_tokens_with_header(tokens: &mut PeekableTokens<T>) -> Result<Self, &'static str> {
+        let dim = infer_geom_dimension(tokens)?;
+        FromTokens::from_tokens_with_parens(tokens, dim)
+    }
+
+    fn from_tokens_with_parens(
+        tokens: &mut PeekableTokens<T>,
+        dim: Dimension,
+    ) -> Result<Self, &'static str> {
         match tokens.next().transpose()? {
             Some(Token::ParenOpen) => (),
             Some(Token::Word(ref s)) if s.eq_ignore_ascii_case("EMPTY") => {
-                return Ok(Default::default())
+                // TODO: expand this to support Z EMPTY
+                // Maybe create a DefaultXY, DefaultXYZ trait etc for each geometry type, and then
+                // here match on the dim to decide which default trait to use.
+                return Ok(Default::default());
             }
             _ => return Err("Missing open parenthesis for type"),
         };
-        let result = FromTokens::from_tokens(tokens);
+        let result = FromTokens::from_tokens(tokens, dim);
         match tokens.next().transpose()? {
             Some(Token::ParenClose) => (),
             _ => return Err("Missing closing parenthesis for type"),
@@ -317,26 +331,31 @@ where
 
     fn from_tokens_with_optional_parens(
         tokens: &mut PeekableTokens<T>,
+        dim: Dimension,
     ) -> Result<Self, &'static str> {
         match tokens.peek() {
-            Some(Ok(Token::ParenOpen)) => Self::from_tokens_with_parens(tokens),
-            _ => Self::from_tokens(tokens),
+            Some(Ok(Token::ParenOpen)) => Self::from_tokens_with_parens(tokens, dim),
+            _ => Self::from_tokens(tokens, dim),
         }
     }
 
-    fn comma_many<F>(f: F, tokens: &mut PeekableTokens<T>) -> Result<Vec<Self>, &'static str>
+    fn comma_many<F>(
+        f: F,
+        tokens: &mut PeekableTokens<T>,
+        dim: Dimension,
+    ) -> Result<Vec<Self>, &'static str>
     where
-        F: Fn(&mut PeekableTokens<T>) -> Result<Self, &'static str>,
+        F: Fn(&mut PeekableTokens<T>, Dimension) -> Result<Self, &'static str>,
     {
         let mut items = Vec::new();
 
-        let item = f(tokens)?;
+        let item = f(tokens, dim)?;
         items.push(item);
 
         while let Some(&Ok(Token::Comma)) = tokens.peek() {
             tokens.next(); // throw away comma
 
-            let item = f(tokens)?;
+            let item = f(tokens, dim)?;
             items.push(item);
         }
 
@@ -404,7 +423,7 @@ mod tests {
         }
 
         // point(x, y, z)
-        let wkt = <Wkt<f64>>::from_str("POINTZ (10 20.1 5)").ok().unwrap();
+        let wkt = <Wkt<f64>>::from_str("POINT Z (10 20.1 5)").ok().unwrap();
         match wkt.item {
             Geometry::Point(Point(Some(coord))) => {
                 assert_eq!(coord.x, 10.0);
@@ -416,7 +435,7 @@ mod tests {
         }
 
         // point(x, y, m)
-        let wkt = <Wkt<f64>>::from_str("POINTM (10 20.1 80)").ok().unwrap();
+        let wkt = <Wkt<f64>>::from_str("POINT M (10 20.1 80)").ok().unwrap();
         match wkt.item {
             Geometry::Point(Point(Some(coord))) => {
                 assert_eq!(coord.x, 10.0);
@@ -428,7 +447,9 @@ mod tests {
         }
 
         // point(x, y, z, m)
-        let wkt = <Wkt<f64>>::from_str("POINTZM (10 20.1 5 80)").ok().unwrap();
+        let wkt = <Wkt<f64>>::from_str("POINT ZM (10 20.1 5 80)")
+            .ok()
+            .unwrap();
         match wkt.item {
             Geometry::Point(Point(Some(coord))) => {
                 assert_eq!(coord.x, 10.0);

diff --git a/src/types/coord.rs b/src/types/coord.rs
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use crate::tokenizer::{PeekableTokens, Token};
+use crate::types::Dimension;
 use crate::{FromTokens, WktNum};
 use std::fmt;
 use std::str::FromStr;
@@ -48,7 +49,7 @@ impl<T> FromTokens<T> for Coord<T>
 where
     T: WktNum + FromStr + Default,
 {
-    fn from_tokens(tokens: &mut PeekableTokens<T>) -> Result<Self, &'static str> {
+    fn from_tokens(tokens: &mut PeekableTokens<T>, dim: Dimension) -> Result<Self, &'static str> {
         let x = match tokens.next().transpose()? {
             Some(Token::Number(n)) => n,
             _ => return Err("Expected a number for the X coordinate"),
@@ -61,17 +62,33 @@ where
         let mut z = None;
         let mut m = None;
 
-        if let Some(Ok(Token::Number(_))) = tokens.peek() {
-            z = match tokens.next().transpose()? {
-                Some(Token::Number(n)) => Some(n),
-                _ => None,
-            };
-
-            if let Some(Ok(Token::Number(_))) = tokens.peek() {
-                m = match tokens.next().transpose()? {
-                    Some(Token::Number(n)) => Some(n),
-                    _ => None,
-                };
+        match dim {
+            Dimension::XY => (),
+            Dimension::XYZ => match tokens.next().transpose()? {
+                Some(Token::Number(n)) => {
+                    z = Some(n);
+                }
+                _ => return Err("Expected a number for the Z coordinate"),
+            },
+            Dimension::XYM => match tokens.next().transpose()? {
+                Some(Token::Number(n)) => {
+                    m = Some(n);
+                }
+                _ => return Err("Expected a number for the M coordinate"),
+            },
+            Dimension::XYZM => {
+                match tokens.next().transpose()? {
+                    Some(Token::Number(n)) => {
+                        z = Some(n);
+                    }
+                    _ => return Err("Expected a number for the Z coordinate"),
+                }
+                match tokens.next().transpose()? {
+                    Some(Token::Number(n)) => {
+                        m = Some(n);
+                    }
+                    _ => return Err("Expected a number for the M coordinate"),
+                }
             }
         }
 

diff --git a/src/types/dimension.rs b/src/types/dimension.rs
@@ -0,0 +1,10 @@
+/// The dimension of geometry that we're parsing.
+#[allow(clippy::upper_case_acronyms)]
+#[derive(Clone, Copy, Debug, Default)]
+pub enum Dimension {
+    #[default]
+    XY,
+    XYZ,
+    XYM,
+    XYZM,
+}
diff --git a/src/types/geometrycollection.rs b/src/types/geometrycollection.rs
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use crate::tokenizer::{PeekableTokens, Token};
+use crate::types::Dimension;
 use crate::{FromTokens, Geometry, WktNum};
 use std::fmt;
 use std::str::FromStr;
@@ -53,7 +54,10 @@ impl<T> FromTokens<T> for GeometryCollection<T>
 where
     T: WktNum + FromStr + Default,
 {
-    fn from_tokens(tokens: &mut PeekableTokens<T>) -> Result<Self, &'static str> {
+    // Unsure if the dimension should be used in parsing GeometryCollection; is it
+    // GEOMETRYCOLLECTION ( POINT Z (...) , POINT ZM (...))
+    // or does a geometry collection have a known dimension?
+    fn from_tokens(tokens: &mut PeekableTokens<T>, _dim: Dimension) -> Result<Self, &'static str> {
         let mut items = Vec::new();
 
         let word = match tokens.next().transpose()? {

diff --git a/src/types/linestring.rs b/src/types/linestring.rs
@@ -14,6 +14,7 @@
 
 use crate::tokenizer::PeekableTokens;
 use crate::types::coord::Coord;
+use crate::types::Dimension;
 use crate::{FromTokens, Geometry, WktNum};
 use std::fmt;
 use std::str::FromStr;
@@ -34,8 +35,8 @@ impl<T> FromTokens<T> for LineString<T>
 where
     T: WktNum + FromStr + Default,
 {
-    fn from_tokens(tokens: &mut PeekableTokens<T>) -> Result<Self, &'static str> {
-        let result = FromTokens::comma_many(<Coord<T> as FromTokens<T>>::from_tokens, tokens);
+    fn from_tokens(tokens: &mut PeekableTokens<T>, dim: Dimension) -> Result<Self, &'static str> {
+        let result = FromTokens::comma_many(<Coord<T> as FromTokens<T>>::from_tokens, tokens, dim);
         result.map(LineString)
     }
 }

diff --git a/src/types/mod.rs b/src/types/mod.rs
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 pub use self::coord::Coord;
+pub use self::dimension::Dimension;
 pub use self::geometrycollection::GeometryCollection;
 pub use self::linestring::LineString;
 pub use self::multilinestring::MultiLineString;
@@ -22,6 +23,7 @@ pub use self::point::Point;
 pub use self::polygon::Polygon;
 
 mod coord;
+mod dimension;
 mod geometrycollection;
 mod linestring;
 mod multilinestring;

diff --git a/src/types/multilinestring.rs b/src/types/multilinestring.rs
@@ -14,6 +14,7 @@
 
 use crate::tokenizer::PeekableTokens;
 use crate::types::linestring::LineString;
+use crate::types::Dimension;
 use crate::{FromTokens, Geometry, WktNum};
 use std::fmt;
 use std::str::FromStr;
@@ -59,10 +60,11 @@ impl<T> FromTokens<T> for MultiLineString<T>
 where
     T: WktNum + FromStr + Default,
 {
-    fn from_tokens(tokens: &mut PeekableTokens<T>) -> Result<Self, &'static str> {
+    fn from_tokens(tokens: &mut PeekableTokens<T>, dim: Dimension) -> Result<Self, &'static str> {
         let result = FromTokens::comma_many(
             <LineString<T> as FromTokens<T>>::from_tokens_with_parens,
             tokens,
+            dim,
         );
         result.map(MultiLineString)
     }

diff --git a/src/types/multipoint.rs b/src/types/multipoint.rs
@@ -14,6 +14,7 @@
 
 use crate::tokenizer::PeekableTokens;
 use crate::types::point::Point;
+use crate::types::Dimension;
 use crate::{FromTokens, Geometry, WktNum};
 use std::fmt;
 use std::str::FromStr;
@@ -55,10 +56,11 @@ impl<T> FromTokens<T> for MultiPoint<T>
 where
     T: WktNum + FromStr + Default,
 {
-    fn from_tokens(tokens: &mut PeekableTokens<T>) -> Result<Self, &'static str> {
+    fn from_tokens(tokens: &mut PeekableTokens<T>, dim: Dimension) -> Result<Self, &'static str> {
         let result = FromTokens::comma_many(
             <Point<T> as FromTokens<T>>::from_tokens_with_optional_parens,
             tokens,
+            dim,
         );
         result.map(MultiPoint)
     }