improve formatting

TabbyML · Aug 26, 2024 · 4d1d19a · 4d1d19a
1 parent 538ffa4
commit 4d1d19a
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/tabby-index-cli/Cargo.toml b/crates/tabby-index-cli/Cargo.toml
@@ -10,4 +10,5 @@ clap = { workspace = true, features = ["derive"]}
 tantivy.workspace = true
 tabby-common = { path = "../tabby-common" }
 anyhow.workspace = true
-serde.workspace = true
+serde.workspace = true
+serde_json.workspace = true
diff --git a/crates/tabby-index-cli/src/commands/head.rs b/crates/tabby-index-cli/src/commands/head.rs
@@ -40,9 +40,9 @@ pub fn run_head_cli(index_path: &Path, args: &HeadArgs) -> anyhow::Result<()> {
                 let doc: TantivyDocument =
                     searcher.doc(doc_address).expect("Failed to read document");
 
-                let json = doc.to_json(&schema.schema);
+                let json_value = to_json_value(doc, &schema.schema);
 
-                println!("{}", json);
+                println!("{}", json_value);
 
                 count += 1;
                 if count >= args.num_docs {
@@ -55,3 +55,18 @@ pub fn run_head_cli(index_path: &Path, args: &HeadArgs) -> anyhow::Result<()> {
 
     Ok(())
 }
+
+fn to_json_value(doc: TantivyDocument, schema: &tantivy::schema::Schema) -> serde_json::Value {
+    let json = doc.to_json(schema);
+    let mut doc: serde_json::Value = serde_json::from_str(&json).expect("Failed to parse JSON");
+
+    for (_, value) in doc.as_object_mut().expect("Expected object").iter_mut() {
+        if let Some(array) = value.as_array_mut() {
+            if array.len() == 1 {
+                *value = array[0].clone();
+            }
+        }
+    }
+
+    doc
+}