Advertisement
homer512

CSV rust

Nov 2nd, 2017
120
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Rust 2.18 KB | None | 0 0
  1. use std::fs;
  2. use std::io::BufRead;
  3. use std::io::Write;
  4. use std::io;
  5. use std::env;
  6.  
  7.  
  8. fn replace_col(inpath: &str, outpath: &str, column: &str,
  9.                    replacement: &str) -> io::Result<()>
  10. {
  11.     let infile = fs::File::open(inpath)?;
  12.     let reader = io::BufReader::new(infile);
  13.     let outfile = fs::File::create(outpath)?;
  14.     let mut writer = io::BufWriter::new(outfile);
  15.     let mut replacenum: usize = 0;
  16.     let mut lines = reader.lines();
  17.     // handle header
  18.     if let Some(line) = lines.next() { // could be empty file
  19.         let header = line?; // could be IO error
  20.         let mut splitted = header.split(',');
  21.         // search for column label
  22.         if let Some(colnum) = splitted.position(|colname| colname == column) {
  23.             replacenum = colnum;
  24.         }
  25.         else {
  26.             return Err(io::Error::new(
  27.                 io::ErrorKind::InvalidData, "Column not found"));
  28.         }
  29.         // pass header through unchanged
  30.         writer.write_fmt(format_args!("{}\n", header))?;
  31.     }
  32.     for line in lines {
  33.         let content = line?; // could be IO error
  34.         let splitted = content.splitn(replacenum + 2, ',');
  35.         // figure out where to insert the replacement string
  36.         let mut replaced = splitted.enumerate()
  37.             .map(|(colnum, orig)|
  38.                  if colnum == replacenum { replacement }
  39.                  else { orig });
  40.         // first column
  41.         if let Some(value) = replaced.next() { // could be empty string
  42.             writer.write(value.as_bytes())?;
  43.         }
  44.         // all following columns. Include comma separator
  45.         for value in replaced {
  46.             writer.write_fmt(format_args!(",{}", value))?;
  47.         }
  48.         writer.write(b"\n")?;
  49.     }
  50.     writer.flush()?;
  51.     Ok({})
  52. }
  53.  
  54. fn main()
  55. {
  56.     let mut args = env::args().skip(1); // skip program name
  57.     let inpath = args.next().expect("Need input file path");
  58.     let column = args.next().expect("Need column label");
  59.     let replacement = args.next().expect("Need replacement value");
  60.     let outpath = args.next().expect("Need output file path");
  61.     replace_col(&inpath, &outpath, &column, &replacement).expect("IO Error");
  62. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement