Support importing csv files.

This commit is contained in:
chriseth 2022-10-28 22:39:47 +02:00
parent 523147f5df
commit 03f0ff8d38
2 changed files with 62 additions and 3 deletions

View file

@ -1,7 +1,7 @@
use std::collections::HashMap;
use std::fs::File;
use std::io;
use std::io::Read;
use std::io::{self, BufRead};
use std::{collections::HashMap, io::BufReader};
use crate::types::{Address, Edge, U256};
@ -11,6 +11,36 @@ pub fn read_edges_binary(path: &String) -> Result<HashMap<Address, Vec<Edge>>, i
read_edges(&mut f, &address_index)
}
pub fn read_edges_csv(path: &String) -> Result<HashMap<Address, Vec<Edge>>, io::Error> {
let mut result = HashMap::<Address, Vec<Edge>>::new();
let f = BufReader::new(File::open(path)?);
for line in f.lines() {
let line = line?;
match &line.split(',').collect::<Vec<_>>()[..] {
[] => continue,
[from, to, token, capacity] => {
let from = Address::from(unescape(from));
let to = Address::from(unescape(to));
let token = Address::from(unescape(token));
let capacity = U256::from(unescape(capacity));
result.entry(from).or_default().push(Edge {
from,
to,
token,
capacity,
});
}
_ => {
return Result::Err(io::Error::new(
io::ErrorKind::Other,
format!("Expected from,to,token,capacity, but got {line}"),
))
}
}
}
Ok(result)
}
fn read_address_index(file: &mut File) -> Result<HashMap<u32, Address>, io::Error> {
let address_count = read_u32(file)?;
let mut addresses = HashMap::new();
@ -71,3 +101,13 @@ fn read_edges(
}
Ok(edges)
}
fn unescape(input: &str) -> &str {
match input.chars().next() {
Some('"') | Some('\'') => {
assert!(input.len() >= 2 && input.chars().last() == input.chars().next());
&input[1..input.len() - 1]
}
_ => input,
}
}

View file

@ -1,5 +1,5 @@
use crate::graph;
use crate::io::read_edges_binary;
use crate::io::{read_edges_binary, read_edges_csv};
use crate::types::{Address, Edge, U256};
use json::JsonValue;
use std::collections::HashMap;
@ -70,6 +70,15 @@ fn handle_connection(
};
socket.write_all(response.as_bytes())?;
}
"load_edges_csv" => {
let response = match load_edges_csv(edges, &request.params["file"].to_string()) {
Ok(len) => jsonrpc_response(request.id, len),
Err(e) => {
jsonrpc_error_response(request.id, -32000, &format!("Error loading edges: {e}"))
}
};
socket.write_all(response.as_bytes())?;
}
"compute_transfer" => {
println!("Computing flow");
let e = edges.read().unwrap().clone();
@ -91,6 +100,16 @@ fn load_edges_binary(
Ok(len)
}
fn load_edges_csv(
edges: &RwLock<Arc<HashMap<Address, Vec<Edge>>>>,
file: &String,
) -> Result<usize, Box<dyn Error>> {
let updated_edges = read_edges_csv(file)?;
let len = updated_edges.len();
*edges.write().unwrap() = Arc::new(updated_edges);
Ok(len)
}
fn compute_transfer(
request: JsonRpcRequest,
edges: &HashMap<Address, Vec<Edge>>,