From adfdae0e64d2e8551a65532f697e8a855c497491 Mon Sep 17 00:00:00 2001
From: Chris Tam <ohgodtamit@gmail.com>
Date: Tue, 21 Nov 2023 14:44:21 -0500
Subject: [PATCH] Support passing dict to add JSON field (#158)

---
 src/document.rs       | 25 +++++++++++++++++--------
 tests/tantivy_test.py |  8 ++++----
 2 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/src/document.rs b/src/document.rs
index ae640d0..8768463 100644
--- a/src/document.rs
+++ b/src/document.rs
@@ -674,19 +674,28 @@ impl Document {
         self.add_value(field_name, bytes);
     }
 
-    /// Add a bytes value to the document.
+    /// Add a JSON value to the document.
     ///
     /// Args:
     ///     field_name (str): The field for which we are adding the bytes.
-    ///     value (str): The json object that will be added to the document.
+    ///     value (str | Dict[str, Any]): The JSON object that will be added
+    ///         to the document.
     ///
-    /// Raises a ValueError if the json is invalid.
-    fn add_json(&mut self, field_name: String, json: &str) -> PyResult<()> {
-        let json_object: serde_json::Value =
-            serde_json::from_str(json).map_err(to_pyerr)?;
-        self.add_value(field_name, json_object);
+    /// Raises a ValueError if the JSON is invalid.
+    fn add_json(&mut self, field_name: String, value: &PyAny) -> PyResult<()> {
+        type JsonMap = serde_json::Map<String, serde_json::Value>;
 
-        Ok(())
+        if let Ok(json_str) = value.extract::<&str>() {
+            let json_map: JsonMap =
+                serde_json::from_str(json_str).map_err(to_pyerr)?;
+            self.add_value(field_name, json_map);
+            Ok(())
+        } else if let Ok(json_map) = pythonize::depythonize::<JsonMap>(value) {
+            self.add_value(field_name, json_map);
+            Ok(())
+        } else {
+            Err(to_pyerr("Invalid JSON object. Expected valid JSON string or Dict[str, Any]."))
+        }
     }
 
     /// Returns the number of added fields that have been added to the document
diff --git a/tests/tantivy_test.py b/tests/tantivy_test.py
index 5e40920..6b4c2b4 100644
--- a/tests/tantivy_test.py
+++ b/tests/tantivy_test.py
@@ -786,12 +786,12 @@ class TestJsonField:
         doc = Document()
         doc.add_json(
             "attributes",
-            """{
-                "order":1.2,
+            {
+                "order": 1.2,
                 "target": "submit-button",
                 "cart": {"product_id": 133},
-                "description": "das keyboard"
-            }""",
+                "description": "das keyboard",
+            },
         )
 
         writer.add_document(doc)