blazegraph_io_core/graphs/
serialization.rs1use crate::types::*;
2use anyhow::Result;
3
4impl DocumentGraph {
5 pub fn to_sequential_format(&self) -> SequentialDocument {
6 let mut nodes: Vec<&DocumentNode> = self.nodes.values().collect();
8 nodes.sort_by(|a, b| {
9 match (a.text_order, b.text_order) {
11 (None, None) => std::cmp::Ordering::Equal,
12 (None, Some(_)) => std::cmp::Ordering::Less,
13 (Some(_), None) => std::cmp::Ordering::Greater,
14 (Some(a_order), Some(b_order)) => a_order.cmp(&b_order),
15 }
16 });
17
18 let segments: Vec<SequentialSegment> = nodes
19 .into_iter()
20 .enumerate()
21 .map(|(index, node)| SequentialSegment {
22 id: index,
23 node_type: node.node_type.clone(),
24 text: node.content.text.clone(),
25 location: node.location.clone(),
26 style: node.style_info.clone(),
27 tokens: node.token_count,
28 })
29 .collect();
30
31 SequentialDocument {
32 format: "sequential".to_string(),
33 segments,
34 structural_profile: self.structural_profile.clone(),
35 }
36 }
37
38 pub fn to_flat_format(&self) -> FlatDocument {
39 let mut nodes: Vec<&DocumentNode> = self.nodes.values().collect();
41 nodes.sort_by(|a, b| {
42 match (a.text_order, b.text_order) {
44 (None, None) => std::cmp::Ordering::Equal,
45 (None, Some(_)) => std::cmp::Ordering::Less,
46 (Some(_), None) => std::cmp::Ordering::Greater,
47 (Some(a_order), Some(b_order)) => a_order.cmp(&b_order),
48 }
49 });
50
51 let chunks: Vec<String> = nodes
52 .into_iter()
53 .map(|node| node.content.text.clone())
54 .collect();
55
56 FlatDocument {
57 format: "flat".to_string(),
58 chunks,
59 }
60 }
61
62 pub fn save_with_format(&self, path: &str, format: &str) -> Result<()> {
63 match format {
64 "sequential" => {
65 let sequential = self.to_sequential_format();
66 let json = serde_json::to_string_pretty(&sequential)?;
67 std::fs::write(path, json)?;
68 }
69 "flat" => {
70 let flat = self.to_flat_format();
71 let json = serde_json::to_string_pretty(&flat)?;
72 std::fs::write(path, json)?;
73 }
74 _ => {
75 self.save_to_json(path)?;
76 }
77 }
78 Ok(())
79 }
80}