airframe-json: DSL for JSON extraction (#538)

wvlet · Jul 7, 2019 · 21a6204 · 21a6204
1 parent a031ea7
commit 21a6204
Show file tree

Hide file tree

Showing 3 changed files with 87 additions and 1 deletion.
diff --git a/airframe-json/src/main/scala/wvlet/airframe/json/JSON.scala b/airframe-json/src/main/scala/wvlet/airframe/json/JSON.scala
@@ -118,6 +118,12 @@ object JSON extends LogSupport {
       s.append("}")
       s.result()
     }
+    def get(name: String): Option[JSONValue] = {
+      v.collectFirst {
+        case (key, value) if key == name =>
+          value
+      }
+    }
   }
   final case class JSONArray(v: IndexedSeq[JSONValue]) extends JSONValue {
     override def toJSON: String = {
@@ -127,6 +133,10 @@ object JSON extends LogSupport {
       s.append("]")
       s.result()
     }
+
+    def apply(i: Int): JSONValue = {
+      v.apply(i)
+    }
   }
 
   /**

diff --git a/airframe-json/src/main/scala/wvlet/airframe/json/package.scala b/airframe-json/src/main/scala/wvlet/airframe/json/package.scala
@@ -12,7 +12,16 @@
  * limitations under the License.
  */
 package wvlet.airframe
-import wvlet.airframe.json.JSON.JSONValue
+import wvlet.airframe.json.JSON.{
+  JSONArray,
+  JSONBoolean,
+  JSONDouble,
+  JSONLong,
+  JSONNull,
+  JSONObject,
+  JSONString,
+  JSONValue
+}
 
 /**
   *
@@ -24,4 +33,48 @@ package object json {
   implicit class RichJson(val json: Json) extends AnyVal {
     def toJSONValue: JSONValue = JSON.parseAny(json)
   }
+
+  implicit class JSONValueOps(val jsonValue: JSONValue) extends AnyVal {
+    def /(name: String): Seq[JSONValue] = {
+      jsonValue match {
+        case jsonObject: JSONObject =>
+          jsonObject.v.collect {
+            case (key, value) if key == name =>
+              value
+          }
+        case jsonArray: JSONArray =>
+          jsonArray.v.flatMap {
+            case value =>
+              value / name
+          }
+        case _ => Nil
+      }
+    }
+
+    def value: Any = {
+      jsonValue match {
+        case JSONNull       => null
+        case JSONDouble(x)  => x
+        case JSONLong(x)    => x
+        case JSONString(x)  => x
+        case JSONBoolean(x) => x
+        case JSONArray(x)   => x.map(_.value)
+        case JSONObject(x)  => x.map(x => (x._1, x._2.value)).toMap
+      }
+    }
+  }
+
+  implicit class JSONValueSeqOps(val jsonValues: Seq[JSONValue]) extends AnyVal {
+    def /(name: String): Seq[JSONValue] = {
+      jsonValues.flatMap { jsonValue =>
+        jsonValue / name
+      }
+    }
+
+    def values: Seq[Any] = {
+      jsonValues.map { jsonValue =>
+        jsonValue.value
+      }
+    }
+  }
 }
diff --git a/airframe-json/src/test/scala/wvlet/airframe/json/JSONTest.scala b/airframe-json/src/test/scala/wvlet/airframe/json/JSONTest.scala
@@ -13,6 +13,7 @@
  */
 package wvlet.airframe.json
 import wvlet.airframe.AirframeSpec
+import wvlet.airframe.json.JSON.{JSONArray, JSONLong, JSONNumber, JSONObject}
 
 /**
   *
@@ -22,4 +23,26 @@ class JSONTest extends AirframeSpec {
     val json: Json = """{"id":1}"""
     json.toJSONValue shouldBe JSON.parse(json)
   }
+
+  "JSONObject.get() and JSONArray.apply()" in {
+    val json: Json = """{"user": [{ "id": 1 }, { "id": 2 }]}"""
+    val jsonValue  = JSON.parse(json)
+
+    val id = for {
+      users <- jsonValue.asInstanceOf[JSONObject].get("user")
+      user  <- Some(users.asInstanceOf[JSONArray](0))
+      id    <- user.asInstanceOf[JSONObject].get("id")
+    } yield id.asInstanceOf[JSONLong].v
+
+    id shouldBe Some(1)
+  }
+
+  "JSON DSL" in {
+    val json: Json = """{"user": [{ "id": 1 }, { "id": 2 }]}"""
+    val jsonValue  = JSON.parse(json)
+
+    val id = (jsonValue / "user" / "id")(0).value
+
+    id shouldBe 1
+  }
 }