From 021f57b3232caa5a469b944ba8757a8d24eeb9e2 Mon Sep 17 00:00:00 2001
From: vms <michail.vms@gmail.com>
Date: Wed, 23 Dec 2020 15:12:58 +0300
Subject: [PATCH] support flatenning

---
 Cargo.toml        |  2 +-
 src/select/mod.rs | 36 +++++++++++++++++++++++++++++++++++-
 tests/selector.rs | 13 +++++++++++++
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index ae6997d..959d02c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "jsonpath_lib-fl"
-version = "0.2.5"
+version = "0.2.6"
 authors = ["Changseok Han <freestrings@gmail.com>"]
 
 description = "It is JsonPath engine written in Rust. it provide a similar API interface in Webassembly and Javascript too. - Webassembly Demo: https://freestrings.github.io/jsonpath"
diff --git a/src/select/mod.rs b/src/select/mod.rs
index 20542ef..7c593f0 100644
--- a/src/select/mod.rs
+++ b/src/select/mod.rs
@@ -40,6 +40,7 @@ enum FilterKey {
 pub enum JsonPathError {
     EmptyPath,
     EmptyValue,
+    CantFlatten(Vec<Value>),
     Path(String),
     Serde(String),
 }
@@ -55,6 +56,9 @@ impl fmt::Display for JsonPathError {
         match self {
             JsonPathError::EmptyPath => f.write_str("path not set"),
             JsonPathError::EmptyValue => f.write_str("json value not set"),
+            JsonPathError::CantFlatten(values) => {
+                write!(f, "json value '{:?}' can't be flattened", values)
+            }
             JsonPathError::Path(msg) => f.write_str(&format!("path error: \n{}\n", msg)),
             JsonPathError::Serde(msg) => f.write_str(&format!("serde error: \n{}\n", msg)),
         }
@@ -455,6 +459,7 @@ pub struct Selector<'a, 'b> {
     chose_indices: Vec<usize>,
     selectors: Vec<Selector<'a, 'b>>,
     selector_filter: FilterTerms<'a>,
+    should_flatten_arrays: bool,
 }
 
 impl<'a, 'b> Selector<'a, 'b> {
@@ -464,6 +469,15 @@ impl<'a, 'b> Selector<'a, 'b> {
 
     pub fn str_path(&mut self, path: &str) -> Result<&mut Self, JsonPathError> {
         debug!("path : {}", path);
+
+        let path = match path.strip_suffix('!') {
+            Some(path) => {
+                self.should_flatten_arrays = true;
+                path
+            }
+            None => path,
+        };
+
         self.node_ref.take();
         self.node = Some(Parser::compile(path).map_err(JsonPathError::Path)?);
         Ok(self)
@@ -556,7 +570,27 @@ impl<'a, 'b> Selector<'a, 'b> {
         self._select()?;
 
         match &self.current {
-            Some(r) => Ok(r.to_vec()),
+            Some(r) => {
+                if self.should_flatten_arrays {
+                    if r.len() != 1 {
+                        let value = r.iter().map(|&v| v.clone()).collect::<Vec<_>>();
+                        return Err(JsonPathError::CantFlatten(value));
+                    }
+                    let value = r[0];
+
+                    return match value {
+                        Value::Array(array) => {
+                            let result: Vec<&Value> = array.iter().map(|v| v).collect::<Vec<_>>();
+                            Ok(result)
+                        }
+                        _ => {
+                            let value = r.iter().map(|&v| v.clone()).collect::<Vec<_>>();
+                            Err(JsonPathError::CantFlatten(value))
+                        }
+                    };
+                }
+                Ok(r.to_vec())
+            }
             _ => Err(JsonPathError::EmptyValue),
         }
     }
diff --git a/tests/selector.rs b/tests/selector.rs
index 975a277..be0d12b 100644
--- a/tests/selector.rs
+++ b/tests/selector.rs
@@ -175,3 +175,16 @@ fn iter_test() {
 
     assert_eq!(result, vec![&json!([1, 2, 3, 4])]);
 }
+
+#[test]
+fn flattening_test() {
+    let array = vec![1, 2, 3, 4, 5];
+    let json_array = json!(array);
+    let haystack = json!([json_array]);
+    let result = jsonpath::select(&haystack, "$.[0]!").unwrap();
+    assert_eq!(result, array);
+
+    let haystack = json!({ "array": array });
+    let result = jsonpath::select(&haystack, "$.array!").unwrap();
+    assert_eq!(result, array);
+}