Skip to main content

codegraph/export/
csv.rs

1//! CSV format export for data analysis in spreadsheets and pandas.
2//!
3//! Generates separate CSV files for nodes and edges with auto-detected columns.
4
5use crate::{CodeGraph, Result};
6use std::collections::HashSet;
7use std::fs::File;
8use std::io::Write;
9use std::path::Path;
10
11/// Export nodes to CSV file
12pub fn export_csv_nodes(graph: &CodeGraph, path: &Path) -> Result<()> {
13    let mut file = File::create(path).map_err(|e| crate::GraphError::Storage {
14        message: format!("Failed to create CSV file: {}", path.display()),
15        source: Some(Box::new(e)),
16    })?;
17
18    // Collect all property keys used in the graph
19    let mut all_keys = HashSet::new();
20    for node_id in 0..graph.node_count() as u64 {
21        if let Ok(node) = graph.get_node(node_id) {
22            for (key, _) in node.properties.iter() {
23                all_keys.insert(key.clone());
24            }
25        }
26    }
27
28    let mut keys_vec: Vec<String> = all_keys.into_iter().collect();
29    keys_vec.sort();
30
31    // Write header
32    write!(file, "id,type").map_err(|e| crate::GraphError::Storage {
33        message: "Failed to write CSV header".to_string(),
34        source: Some(Box::new(e)),
35    })?;
36    for key in &keys_vec {
37        write!(file, ",{key}").map_err(|e| crate::GraphError::Storage {
38            message: "Failed to write CSV header".to_string(),
39            source: Some(Box::new(e)),
40        })?;
41    }
42    writeln!(file).map_err(|e| crate::GraphError::Storage {
43        message: "Failed to write CSV header".to_string(),
44        source: Some(Box::new(e)),
45    })?;
46
47    // Write rows
48    for node_id in 0..graph.node_count() as u64 {
49        if let Ok(node) = graph.get_node(node_id) {
50            write!(file, "{},{:?}", node_id, node.node_type).map_err(|e| {
51                crate::GraphError::Storage {
52                    message: "Failed to write CSV row".to_string(),
53                    source: Some(Box::new(e)),
54                }
55            })?;
56
57            for key in &keys_vec {
58                write!(file, ",").map_err(|e| crate::GraphError::Storage {
59                    message: "Failed to write CSV row".to_string(),
60                    source: Some(Box::new(e)),
61                })?;
62                if let Some(value) = node.properties.get(key) {
63                    write!(file, "{}", escape_csv(&format_property_value(value))).map_err(|e| {
64                        crate::GraphError::Storage {
65                            message: "Failed to write CSV row".to_string(),
66                            source: Some(Box::new(e)),
67                        }
68                    })?;
69                }
70            }
71            writeln!(file).map_err(|e| crate::GraphError::Storage {
72                message: "Failed to write CSV row".to_string(),
73                source: Some(Box::new(e)),
74            })?;
75        }
76    }
77
78    Ok(())
79}
80
81/// Export edges to CSV file
82pub fn export_csv_edges(graph: &CodeGraph, path: &Path) -> Result<()> {
83    let mut file = File::create(path).map_err(|e| crate::GraphError::Storage {
84        message: format!("Failed to create CSV file: {}", path.display()),
85        source: Some(Box::new(e)),
86    })?;
87
88    // Collect all property keys used in edges
89    let mut all_keys = HashSet::new();
90    for edge_id in 0..graph.edge_count() as u64 {
91        if let Ok(edge) = graph.get_edge(edge_id) {
92            for (key, _) in edge.properties.iter() {
93                all_keys.insert(key.clone());
94            }
95        }
96    }
97
98    let mut keys_vec: Vec<String> = all_keys.into_iter().collect();
99    keys_vec.sort();
100
101    // Write header
102    write!(file, "id,source,target,type").map_err(|e| crate::GraphError::Storage {
103        message: "Failed to write CSV header".to_string(),
104        source: Some(Box::new(e)),
105    })?;
106    for key in &keys_vec {
107        write!(file, ",{key}").map_err(|e| crate::GraphError::Storage {
108            message: "Failed to write CSV header".to_string(),
109            source: Some(Box::new(e)),
110        })?;
111    }
112    writeln!(file).map_err(|e| crate::GraphError::Storage {
113        message: "Failed to write CSV header".to_string(),
114        source: Some(Box::new(e)),
115    })?;
116
117    // Write rows
118    for edge_id in 0..graph.edge_count() as u64 {
119        if let Ok(edge) = graph.get_edge(edge_id) {
120            write!(
121                file,
122                "{},{},{},{:?}",
123                edge_id, edge.source_id, edge.target_id, edge.edge_type
124            )
125            .map_err(|e| crate::GraphError::Storage {
126                message: "Failed to write CSV row".to_string(),
127                source: Some(Box::new(e)),
128            })?;
129
130            for key in &keys_vec {
131                write!(file, ",").map_err(|e| crate::GraphError::Storage {
132                    message: "Failed to write CSV row".to_string(),
133                    source: Some(Box::new(e)),
134                })?;
135                if let Some(value) = edge.properties.get(key) {
136                    write!(file, "{}", escape_csv(&format_property_value(value))).map_err(|e| {
137                        crate::GraphError::Storage {
138                            message: "Failed to write CSV row".to_string(),
139                            source: Some(Box::new(e)),
140                        }
141                    })?;
142                }
143            }
144            writeln!(file).map_err(|e| crate::GraphError::Storage {
145                message: "Failed to write CSV row".to_string(),
146                source: Some(Box::new(e)),
147            })?;
148        }
149    }
150
151    Ok(())
152}
153
154/// Export both nodes and edges to separate CSV files (convenience method)
155pub fn export_csv(graph: &CodeGraph, nodes_path: &Path, edges_path: &Path) -> Result<()> {
156    export_csv_nodes(graph, nodes_path)?;
157    export_csv_edges(graph, edges_path)?;
158    Ok(())
159}
160
161/// Format property value for CSV
162fn format_property_value(value: &crate::PropertyValue) -> String {
163    match value {
164        crate::PropertyValue::String(s) => s.clone(),
165        crate::PropertyValue::Int(i) => i.to_string(),
166        crate::PropertyValue::Float(f) => f.to_string(),
167        crate::PropertyValue::Bool(b) => b.to_string(),
168        crate::PropertyValue::StringList(v) => v.join(";"),
169        crate::PropertyValue::IntList(v) => v
170            .iter()
171            .map(|i| i.to_string())
172            .collect::<Vec<_>>()
173            .join(";"),
174        crate::PropertyValue::Null => String::new(),
175    }
176}
177
178/// Escape CSV value (add quotes if contains comma, quote, or newline)
179fn escape_csv(s: &str) -> String {
180    if s.contains(',') || s.contains('"') || s.contains('\n') {
181        format!("\"{}\"", s.replace('"', "\"\""))
182    } else {
183        s.to_string()
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    #[test]
192    fn test_escape_csv() {
193        assert_eq!(escape_csv("hello"), "hello");
194        assert_eq!(escape_csv("hello,world"), "\"hello,world\"");
195        assert_eq!(escape_csv("say \"hi\""), "\"say \"\"hi\"\"\"");
196    }
197}