Add types to represent JSON values (#112)

Author: andrewheard

Sources/GoogleAI/JSONValue.swift

@@ -0,0 +1,71 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+/// A collection of name-value pairs representing a JSON object.
+///
+/// This may be decoded from, or encoded to, a
+/// [`google.protobuf.Struct`](https://protobuf.dev/reference/protobuf/google.protobuf/#struct).
+public typealias JSONObject = [String: JSONValue]
+
+/// Represents a value in one of JSON's data types.
+///
+/// This may be decoded from, or encoded to, a
+/// [`google.protobuf.Value`](https://protobuf.dev/reference/protobuf/google.protobuf/#value).
+public enum JSONValue {
+  /// A `null` value.
+  case null
+
+  /// A numeric value.
+  case number(Double)
+
+  /// A string value.
+  case string(String)
+
+  /// A boolean value.
+  case bool(Bool)
+
+  /// A JSON object.
+  case object(JSONObject)
+
+  /// An array of `JSONValue`s.
+  case array([JSONValue])
+}
+
+extension JSONValue: Decodable {
+  public init(from decoder: Decoder) throws {
+    let container = try decoder.singleValueContainer()
+    if container.decodeNil() {
+      self = .null
+    } else if let numberValue = try? container.decode(Double.self) {
+      self = .number(numberValue)
+    } else if let stringValue = try? container.decode(String.self) {
+      self = .string(stringValue)
+    } else if let boolValue = try? container.decode(Bool.self) {
+      self = .bool(boolValue)
+    } else if let objectValue = try? container.decode(JSONObject.self) {
+      self = .object(objectValue)
+    } else if let arrayValue = try? container.decode([JSONValue].self) {
+      self = .array(arrayValue)
+    } else {
+      throw DecodingError.dataCorruptedError(
+        in: container,
+        debugDescription: "Failed to decode JSON value."
+      )
+    }
+  }
+}
+
+extension JSONValue: Equatable {}

Tests/GoogleAITests/JSONValueTests.swift

@@ -0,0 +1,85 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+import XCTest
+
+@testable import GoogleGenerativeAI
+
+final class JSONValueTests: XCTestCase {
+  func testDecodeNull() throws {
+    let jsonData = try XCTUnwrap("null".data(using: .utf8))
+
+    let jsonObject = try XCTUnwrap(JSONDecoder().decode(JSONValue.self, from: jsonData))
+
+    XCTAssertEqual(jsonObject, .null)
+  }
+
+  func testDecodeNumber() throws {
+    let expectedNumber = 3.14159
+    let jsonData = try XCTUnwrap("\(expectedNumber)".data(using: .utf8))
+
+    let jsonObject = try XCTUnwrap(JSONDecoder().decode(JSONValue.self, from: jsonData))
+
+    XCTAssertEqual(jsonObject, .number(expectedNumber))
+  }
+
+  func testDecodeString() throws {
+    let expectedString = "hello-world"
+    let jsonData = try XCTUnwrap("\"\(expectedString)\"".data(using: .utf8))
+
+    let jsonObject = try XCTUnwrap(JSONDecoder().decode(JSONValue.self, from: jsonData))
+
+    XCTAssertEqual(jsonObject, .string(expectedString))
+  }
+
+  func testDecodeBool() throws {
+    let expectedBool = true
+    let jsonData = try XCTUnwrap("\(expectedBool)".data(using: .utf8))
+
+    let jsonObject = try XCTUnwrap(JSONDecoder().decode(JSONValue.self, from: jsonData))
+
+    XCTAssertEqual(jsonObject, .bool(expectedBool))
+  }
+
+  func testDecodeObject() throws {
+    let numberKey = "pi"
+    let numberValue = 3.14159
+    let stringKey = "hello"
+    let stringValue = "world"
+    let expectedObject: JSONObject = [
+      numberKey: .number(numberValue),
+      stringKey: .string(stringValue),
+    ]
+    let json = """
+    {
+      "\(numberKey)": \(numberValue),
+      "\(stringKey)": "\(stringValue)"
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let jsonObject = try XCTUnwrap(JSONDecoder().decode(JSONValue.self, from: jsonData))
+
+    XCTAssertEqual(jsonObject, .object(expectedObject))
+  }
+
+  func testDecodeArray() throws {
+    let numberValue = 3.14159
+    let expectedArray: [JSONValue] = [.null, .number(numberValue)]
+    let jsonData = try XCTUnwrap("[ null, \(numberValue) ]".data(using: .utf8))
+
+    let jsonObject = try XCTUnwrap(JSONDecoder().decode(JSONValue.self, from: jsonData))
+
+    XCTAssertEqual(jsonObject, .array(expectedArray))
+  }
+}