use crate::{CodeGraph, Result};
use std::collections::HashSet;
use std::fs::File;
use std::io::Write;
use std::path::Path;
pub fn export_csv_nodes(graph: &CodeGraph, path: &Path) -> Result<()> {
let mut file = File::create(path).map_err(|e| crate::GraphError::Storage {
message: format!("Failed to create CSV file: {}", path.display()),
source: Some(Box::new(e)),
})?;
let mut all_keys = HashSet::new();
for node_id in 0..graph.node_count() as u64 {
if let Ok(node) = graph.get_node(node_id) {
for (key, _) in node.properties.iter() {
all_keys.insert(key.clone());
}
}
}
let mut keys_vec: Vec<String> = all_keys.into_iter().collect();
keys_vec.sort();
write!(file, "id,type").map_err(|e| crate::GraphError::Storage {
message: "Failed to write CSV header".to_string(),
source: Some(Box::new(e)),
})?;
for key in &keys_vec {
write!(file, ",{key}").map_err(|e| crate::GraphError::Storage {
message: "Failed to write CSV header".to_string(),
source: Some(Box::new(e)),
})?;
}
writeln!(file).map_err(|e| crate::GraphError::Storage {
message: "Failed to write CSV header".to_string(),
source: Some(Box::new(e)),
})?;
for node_id in 0..graph.node_count() as u64 {
if let Ok(node) = graph.get_node(node_id) {
write!(file, "{},{:?}", node_id, node.node_type).map_err(|e| {
crate::GraphError::Storage {
message: "Failed to write CSV row".to_string(),
source: Some(Box::new(e)),
}
})?;
for key in &keys_vec {
write!(file, ",").map_err(|e| crate::GraphError::Storage {
message: "Failed to write CSV row".to_string(),
source: Some(Box::new(e)),
})?;
if let Some(value) = node.properties.get(key) {
write!(file, "{}", escape_csv(&format_property_value(value))).map_err(|e| {
crate::GraphError::Storage {
message: "Failed to write CSV row".to_string(),
source: Some(Box::new(e)),
}
})?;
}
}
writeln!(file).map_err(|e| crate::GraphError::Storage {
message: "Failed to write CSV row".to_string(),
source: Some(Box::new(e)),
})?;
}
}
Ok(())
}
pub fn export_csv_edges(graph: &CodeGraph, path: &Path) -> Result<()> {
let mut file = File::create(path).map_err(|e| crate::GraphError::Storage {
message: format!("Failed to create CSV file: {}", path.display()),
source: Some(Box::new(e)),
})?;
let mut all_keys = HashSet::new();
for edge_id in 0..graph.edge_count() as u64 {
if let Ok(edge) = graph.get_edge(edge_id) {
for (key, _) in edge.properties.iter() {
all_keys.insert(key.clone());
}
}
}
let mut keys_vec: Vec<String> = all_keys.into_iter().collect();
keys_vec.sort();
write!(file, "id,source,target,type").map_err(|e| crate::GraphError::Storage {
message: "Failed to write CSV header".to_string(),
source: Some(Box::new(e)),
})?;
for key in &keys_vec {
write!(file, ",{key}").map_err(|e| crate::GraphError::Storage {
message: "Failed to write CSV header".to_string(),
source: Some(Box::new(e)),
})?;
}
writeln!(file).map_err(|e| crate::GraphError::Storage {
message: "Failed to write CSV header".to_string(),
source: Some(Box::new(e)),
})?;
for edge_id in 0..graph.edge_count() as u64 {
if let Ok(edge) = graph.get_edge(edge_id) {
write!(
file,
"{},{},{},{:?}",
edge_id, edge.source_id, edge.target_id, edge.edge_type
)
.map_err(|e| crate::GraphError::Storage {
message: "Failed to write CSV row".to_string(),
source: Some(Box::new(e)),
})?;
for key in &keys_vec {
write!(file, ",").map_err(|e| crate::GraphError::Storage {
message: "Failed to write CSV row".to_string(),
source: Some(Box::new(e)),
})?;
if let Some(value) = edge.properties.get(key) {
write!(file, "{}", escape_csv(&format_property_value(value))).map_err(|e| {
crate::GraphError::Storage {
message: "Failed to write CSV row".to_string(),
source: Some(Box::new(e)),
}
})?;
}
}
writeln!(file).map_err(|e| crate::GraphError::Storage {
message: "Failed to write CSV row".to_string(),
source: Some(Box::new(e)),
})?;
}
}
Ok(())
}
pub fn export_csv(graph: &CodeGraph, nodes_path: &Path, edges_path: &Path) -> Result<()> {
export_csv_nodes(graph, nodes_path)?;
export_csv_edges(graph, edges_path)?;
Ok(())
}
fn format_property_value(value: &crate::PropertyValue) -> String {
match value {
crate::PropertyValue::String(s) => s.clone(),
crate::PropertyValue::Int(i) => i.to_string(),
crate::PropertyValue::Float(f) => f.to_string(),
crate::PropertyValue::Bool(b) => b.to_string(),
crate::PropertyValue::StringList(v) => v.join(";"),
crate::PropertyValue::IntList(v) => v
.iter()
.map(|i| i.to_string())
.collect::<Vec<_>>()
.join(";"),
crate::PropertyValue::Null => String::new(),
}
}
fn escape_csv(s: &str) -> String {
if s.contains(',') || s.contains('"') || s.contains('\n') {
format!("\"{}\"", s.replace('"', "\"\""))
} else {
s.to_string()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_escape_csv() {
assert_eq!(escape_csv("hello"), "hello");
assert_eq!(escape_csv("hello,world"), "\"hello,world\"");
assert_eq!(escape_csv("say \"hi\""), "\"say \"\"hi\"\"\"");
}
}