<commit_before>extern crate image;
extern crate image_base64;
extern crate cgi;

use std::fs::File;
use std::path::Path;
use std::env;
use std::io;//::{self, Write};
use image::GenericImage;

//should be a trait.  I am not sure how to impliment one for only Vec <Vec <bool>> and not Vec <T>
fn boundless_insert(y: i64, x: i64, value: bool, img: &mut Vec<Vec <bool>>) {
	let mut tx = x;
	let mut ty = y;
	for i in 0..img.len(){
		let mut ttx = x;
		while ttx < 0 {
			img[i].insert(0, false);
			ttx += 1;
		}
		assert!(((x < 0) && (ttx == 0)) || (x == ttx), "Expected ttx = zero or ttx = x.  Found {}", ttx);
		while img[i].len() <= ttx as usize {
			img[i].push(false);
		}
		tx = ttx
	}
	let mut row:Vec <bool> = Vec::new();
	for i in 0..img[0].len() {
		row.push(false);
	}
	while ty < 0 {
		img.insert(0, row.clone());
		ty += 1;
	}
	assert!(((y < 0) && (ty == 0)) || (y == ty), "Expected ty = zero or ty = y.  Found {}", ty);
	while img.len() <= ty as usize {
		img.push(row.clone())
	}
	img[ty as usize][tx as usize] = value;
}

struct ImgBlob {
	blob_type: u8, //0: object, 1: line
	top_right: [usize; 2],
	middle: [usize; 2],
	text: String,
	bitmap: Vec <Vec <bool>>
}

impl ImgBlob {
	fn from_top_right(x: usize, y: usize, claim: &mut Vec <Vec <bool>>, img: &Vec <Vec <bool>>) -> Option<ImgBlob>{
		//let mut imgbuf = image::ImageBuffer::new(2550u32, 3300u32);
		let mut right = x;
		let mut top = y; //pretty sure that this is not supposed to change and can be eliminated
		let mut bitmap: Vec <Vec <bool>> = Vec::new();
		let mut queue: Vec <[usize; 2]> = Vec::new();
		claim[y][x] = true;
		//imgbuf.put_pixel(x as u32, y as u32, image::Rgb([0, 255, 0]));
		queue.push([y+1, x]);
		queue.push([y, x+1]);
		bitmap.push(vec![true]);
		while queue.len() > 0 {
			let tempx = queue[0][1];
			let tempy = queue[0][0];
			if (tempy < img.len()) && (tempx  0) && (tempx > 0) {
				if img[tempy][tempx] && !(claim[tempy][tempx]){
					boundless_insert(((tempy as i64)-(top as i64)), ((tempx as i64)-(right as i64)), true, &mut bitmap);
					if tempx < right {
						right = tempx;
					}
					if tempy < top {
						top = tempy;
					}
					//imgbuf.put_pixel(tempx as u32, tempy as u32, image::Rgb([0, 255, 0]));
					claim[tempy][tempx] = true;
					if queue[0][1] > top {
						queue.push([tempy-1, tempx]);
					}
					queue.push([tempy, tempx-1]);
					queue.push([tempy+1, tempx]);
					queue.push([tempy, tempx+1]);
				}/* else if !(claim[tempy][tempx]) {
					imgbuf.put_pixel(tempx as u32, tempy as u32, image::Rgb([255, 0, 0]));
				}*/
			}/* else {
				imgbuf.put_pixel(tempx as u32, tempy as u32, image::Rgb([255, 0, 0]));
			}*/
			queue.remove(0);
		}
		if (bitmap[0].len() + bitmap.len() > 6) && (bitmap.len() > 3) && (bitmap[0].len() > 3) {
			/*let mut imgbuf = image::ImageBuffer::new((bitmap[0].len()) as u32, (bitmap.len()) as u32);
			for (x, y, pixel) in imgbuf.enumerate_pixels_mut() {
				*pixel = image::Luma([if bitmap[y as usize][x as usize] {255} else {0}]);
			}
			let xdim = &*bitmap[0].len().to_string();
			let ydim = &*bitmap.len().to_string();
			let ref mut fout = File::create(&Path::new(&(String::from("outg")+xdim+"x"+ydim+".png"))).unwrap();
			let _ = image::ImageRgb8(imgbuf).save(fout, image::PNG);*/
			println!("Found: {} x {}", bitmap[0].len(), bitmap.len());
			Some(ImgBlob {
				blob_type: 0,
				top_right: [right, top],
				middle: [right+(bitmap[0].len()/2), top+(bitmap.len()/2)],
				text: String::from(""),
				bitmap: bitmap
			})
		} else {
			None
		}
	}
}

struct Section {
	Heading: ImgBlob,
	blobs: Vec <ImgBlob>
}

struct Chapter {
	heading: ImgBlob,
	blobs: Vec <ImgBlob>,
	Sections: Vec

}

fn main() {
	let mut claimed: Vec <Vec <bool>> = Vec::new();
	let mut thresh: Vec <Vec <bool>> = Vec::new();
	//let mut rblobs: Vec <ImgBlob> = Vec::new();
	let mut gblobs: Vec <ImgBlob> = Vec::new();
	//let mut bblobs: Vec <ImgBlob> = Vec::new();
	let img = image::open(&Path::new("target.jpg")).unwrap();
	let mut row:Vec <bool> = Vec::new();
	for x in 0..img.width() {
		row.push(false);
	}
	for y in 0..img.height() {
		claimed.push(row.clone());
		thresh.push(row.clone());
	}
	let rgbimg = img.to_rgb();
	//let mut imgbuf = image::ImageBuffer::new(img.width(), img.height());
	/*for (x, y, pixel) in rgbimg.enumerate_pixels() {
		thresh[y as usize][x as usize] = if (pixel[0] > 140) && (pixel[1] <= 170) && (pixel[2] <= 170) {true} else {false};
		imgbuf.put_pixel(x, y, image::Luma([if (pixel[0] > 140) & (pixel[1] <= 170) & (pixel[2] <= 170) {255} else {0}]));
	}
	let ref mut fout = File::create(&Path::new("outr.png")).unwrap();
	let _ = image::ImageLuma8(imgbuf).save(fout, image::PNG);
	for y in 0..thresh.len() {
		for x in 0..thresh[0].len() {
			if thresh[y][x] {
				match ImgBlob::from_top_right(x, y, &mut claimed, &thresh){
					Some(o) => rblobs.push(o),
					None => {},
				}
			}
		}
	}
	println!("Finished R");*/
	for (x, y, pixel) in rgbimg.enumerate_pixels() {
		thresh[y as usize][x as usize] = if (pixel[1] > 140) && (pixel[0] <= 170) && (pixel[2] <= 170) {true} else {false};
	}
	for y in 0..thresh.len() {
		for x in 0..thresh[0].len() {
			//imgbuf.put_pixel(x as u32, y as u32, image::Luma([if thresh[y][x] {255} else {0}]));
			if thresh[y][x] {
				match ImgBlob::from_top_right(x, y, &mut claimed, &thresh){
					Some(o) => gblobs.push(o),
					None => {},
				}
			}
		}
	}
	/*let ref mut fout = File::create(&Path::new("outg.png")).unwrap();
	let _ = image::ImageLuma8(imgbuf).save(fout, image::PNG);*/
	let mut imgbuf = image::ImageBuffer::new(img.width(), img.height());
	for i in 0..gblobs.len() {
		let xoff = gblobs[i].top_right[0];
		let yoff = gblobs[i].top_right[1];
		for y in 0..gblobs[i].bitmap.len() {
			for x in 0..gblobs[i].bitmap[y].len() {
				imgbuf.put_pixel((x+xoff) as u32, (y+yoff) as u32, image::Luma([if (gblobs[i].bitmap[y][x]) {255} else {0}]));
			}
		}
	}
	let ref mut fout = File::create(&Path::new("reconstructg.png")).unwrap();
	let _ = image::ImageLuma8(imgbuf).save(fout, image::PNG);
	println!("Finished G");
	/*for (x, y, pixel) in rgbimg.enumerate_pixels() {
		thresh[y as usize][x as usize] = if (pixel[2] > 140) && (pixel[1] <= 170) && (pixel[0] <= 170) {true} else {false};
	}
	for y in 0..thresh.len() {
		for x in 0..thresh[0].len() {
			imgbuf.put_pixel(x as u32, y as u32, image::Luma([if thresh[y][x] {255} else {0}]));
			if thresh[y][x] {
				match ImgBlob::from_top_right(x, y, &mut claimed, &thresh){
					Some(o) => bblobs.push(o),
					None => {},
				}
			}
		}
	}
	let ref mut fout = File::create(&Path::new("outb.png")).unwrap();
	let _ = image::ImageLuma8(imgbuf).save(fout, image::PNG);
	println!("Finished B");*/

}
<commit_msg>Basic line detection<commit_after>extern crate image;
extern crate image_base64;
extern crate cgi;

use std::fs::File;
use std::path::Path;
use std::env;
use std::io;//::{self, Write};
use image::GenericImage;

//should be a trait.  I am not sure how to impliment one for only Vec <Vec <bool>> and not Vec <T>
fn boundless_insert(y: i64, x: i64, value: bool, img: &mut Vec<Vec <bool>>) {
	let mut tx = x;
	let mut ty = y;
	for i in 0..img.len(){
		let mut ttx = x;
		while ttx < 0 {
			img[i].insert(0, false);
			ttx += 1;
		}
		assert!(((x < 0) && (ttx == 0)) || (x == ttx), "Expected ttx = zero or ttx = x.  Found {}", ttx);
		while img[i].len() <= ttx as usize {
			img[i].push(false);
		}
		tx = ttx
	}
	let mut row:Vec <bool> = Vec::new();
	for i in 0..img[0].len() {
		row.push(false);
	}
	while ty < 0 {
		img.insert(0, row.clone());
		ty += 1;
	}
	assert!(((y < 0) && (ty == 0)) || (y == ty), "Expected ty = zero or ty = y.  Found {}", ty);
	while img.len() <= ty as usize {
		img.push(row.clone())
	}
	img[ty as usize][tx as usize] = value;
}

struct ImgBlob {
	blob_type: u8, //0: object, 1: line
	top_right: [usize; 2],
	middle: [usize; 2],
	text: String,
	bitmap: Vec <Vec <bool>>
}

impl ImgBlob {
	fn from_top_right(x: usize, y: usize, claim: &mut Vec <Vec <bool>>, img: &Vec <Vec <bool>>) -> Option<ImgBlob>{
		//let mut imgbuf = image::ImageBuffer::new(2550u32, 3300u32);
		let mut right = x;
		let mut top = y; //pretty sure that this is not supposed to change and can be eliminated
		let mut bitmap: Vec <Vec <bool>> = Vec::new();
		let mut queue: Vec <[usize; 2]> = Vec::new();
		claim[y][x] = true;
		//imgbuf.put_pixel(x as u32, y as u32, image::Rgb([0, 255, 0]));
		queue.push([y+1, x]);
		queue.push([y, x+1]);
		bitmap.push(vec![true]);
		while queue.len() > 0 {
			let tempx = queue[0][1];
			let tempy = queue[0][0];
			if (tempy < img.len()) && (tempx  0) && (tempx > 0) {
				if img[tempy][tempx] && !(claim[tempy][tempx]){
					boundless_insert(((tempy as i64)-(top as i64)), ((tempx as i64)-(right as i64)), true, &mut bitmap);
					if tempx < right {
						right = tempx;
					}
					if tempy < top {
						top = tempy;
					}
					//imgbuf.put_pixel(tempx as u32, tempy as u32, image::Rgb([0, 255, 0]));
					claim[tempy][tempx] = true;
					if queue[0][1] > top {
						queue.push([tempy-1, tempx]);
					}
					queue.push([tempy, tempx-1]);
					queue.push([tempy+1, tempx]);
					queue.push([tempy, tempx+1]);
				}/* else if !(claim[tempy][tempx]) {
					imgbuf.put_pixel(tempx as u32, tempy as u32, image::Rgb([255, 0, 0]));
				}*/
			}/* else {
				imgbuf.put_pixel(tempx as u32, tempy as u32, image::Rgb([255, 0, 0]));
			}*/
			queue.remove(0);
		}
		if (bitmap[0].len() + bitmap.len() > 6) && (bitmap.len() > 3) && (bitmap[0].len() > 3) {
			/*let mut imgbuf = image::ImageBuffer::new((bitmap[0].len()) as u32, (bitmap.len()) as u32);
			for (x, y, pixel) in imgbuf.enumerate_pixels_mut() {
				*pixel = image::Luma([if bitmap[y as usize][x as usize] {255} else {0}]);
			}
			let xdim = &*bitmap[0].len().to_string();
			let ydim = &*bitmap.len().to_string();
			let ref mut fout = File::create(&Path::new(&(String::from("outg")+xdim+"x"+ydim+".png"))).unwrap();
			let _ = image::ImageRgb8(imgbuf).save(fout, image::PNG);*/
			println!("Found: {} x {}", bitmap[0].len(), bitmap.len());
			Some(ImgBlob {
				blob_type: if (bitmap[0].len() / bitmap.len()) > 10 {1} else {0},
				top_right: [right, top],
				middle: [right+(bitmap[0].len()/2), top+(bitmap.len()/2)],
				text: String::from(""),
				bitmap: bitmap
			})
		} else {
			None
		}
	}
}

struct Section {
	Heading: ImgBlob,
	blobs: Vec <ImgBlob>
}

struct Chapter {
	heading: ImgBlob,
	blobs: Vec <ImgBlob>,
	Sections: Vec

}

fn main() {
	let mut claimed: Vec <Vec <bool>> = Vec::new();
	let mut thresh: Vec <Vec <bool>> = Vec::new();
	//let mut rblobs: Vec <ImgBlob> = Vec::new();
	let mut gblobs: Vec <ImgBlob> = Vec::new();
	//let mut bblobs: Vec <ImgBlob> = Vec::new();
	let img = image::open(&Path::new("target.jpg")).unwrap();
	let mut row:Vec <bool> = Vec::new();
	for x in 0..img.width() {
		row.push(false);
	}
	for y in 0..img.height() {
		claimed.push(row.clone());
		thresh.push(row.clone());
	}
	let rgbimg = img.to_rgb();
	//let mut imgbuf = image::ImageBuffer::new(img.width(), img.height());
	/*for (x, y, pixel) in rgbimg.enumerate_pixels() {
		thresh[y as usize][x as usize] = if (pixel[0] > 140) && (pixel[1] <= 170) && (pixel[2] <= 170) {true} else {false};
		imgbuf.put_pixel(x, y, image::Luma([if (pixel[0] > 140) & (pixel[1] <= 170) & (pixel[2] <= 170) {255} else {0}]));
	}
	let ref mut fout = File::create(&Path::new("outr.png")).unwrap();
	let _ = image::ImageLuma8(imgbuf).save(fout, image::PNG);
	for y in 0..thresh.len() {
		for x in 0..thresh[0].len() {
			if thresh[y][x] {
				match ImgBlob::from_top_right(x, y, &mut claimed, &thresh){
					Some(o) => rblobs.push(o),
					None => {},
				}
			}
		}
	}
	println!("Finished R");*/
	for (x, y, pixel) in rgbimg.enumerate_pixels() {
		thresh[y as usize][x as usize] = if (pixel[1] > 140) && (pixel[0] <= 170) && (pixel[2] <= 170) {true} else {false};
	}
	for y in 0..thresh.len() {
		for x in 0..thresh[0].len() {
			//imgbuf.put_pixel(x as u32, y as u32, image::Luma([if thresh[y][x] {255} else {0}]));
			if thresh[y][x] {
				match ImgBlob::from_top_right(x, y, &mut claimed, &thresh){
					Some(o) => gblobs.push(o),
					None => {},
				}
			}
		}
	}
	/*let ref mut fout = File::create(&Path::new("outg.png")).unwrap();
	let _ = image::ImageLuma8(imgbuf).save(fout, image::PNG);*/
	let mut imgbuf = image::ImageBuffer::new(img.width(), img.height());
	for i in 0..gblobs.len() {
		let xoff = gblobs[i].top_right[0];
		let yoff = gblobs[i].top_right[1];
		for y in 0..gblobs[i].bitmap.len() {
			for x in 0..gblobs[i].bitmap[y].len() {
				if gblobs[i].blob_type == 0 {
					imgbuf.put_pixel((x+xoff) as u32, (y+yoff) as u32, image::Rgb(if (gblobs[i].bitmap[y][x]) {[0, 255, 0]} else {[0, 0, 0]}));
				} else {
					imgbuf.put_pixel((x+xoff) as u32, (y+yoff) as u32, image::Rgb(if (gblobs[i].bitmap[y][x]) {[255, 0, 0]} else {[0, 0, 0]}));
				}
			}
		}
	}
	let ref mut fout = File::create(&Path::new("reconstructg.png")).unwrap();
	let _ = image::ImageRgb8(imgbuf).save(fout, image::PNG);
	println!("Finished G");
	/*for (x, y, pixel) in rgbimg.enumerate_pixels() {
		thresh[y as usize][x as usize] = if (pixel[2] > 140) && (pixel[1] <= 170) && (pixel[0] <= 170) {true} else {false};
	}
	for y in 0..thresh.len() {
		for x in 0..thresh[0].len() {
			imgbuf.put_pixel(x as u32, y as u32, image::Luma([if thresh[y][x] {255} else {0}]));
			if thresh[y][x] {
				match ImgBlob::from_top_right(x, y, &mut claimed, &thresh){
					Some(o) => bblobs.push(o),
					None => {},
				}
			}
		}
	}
	let ref mut fout = File::create(&Path::new("outb.png")).unwrap();
	let _ = image::ImageLuma8(imgbuf).save(fout, image::PNG);
	println!("Finished B");*/

}
<|endoftext|>


<commit_before>#![allow(unstable, dead_code)]

extern crate getopts;
extern crate libc;

mod ansi;
mod matching;
mod screen;

use std::os;
use std::io;
use std::cmp::min;
use screen::Screen;
use screen::Key::*;

fn main() {
  let input_args = os::args();

  let args = match parse_args(&input_args) {
    Some(args) => args,
    None => {
      os::set_exit_status(1);
      return;
    },
  };

  if args.help { return; }

  let mut screen = Screen::open_screen();

  let choices = read_choices();
  let visible_choices = min(20, screen.height - 1);

  let mut search: String = args.initial_search.clone();
  let start_line = screen.height - visible_choices - 1;
  loop {
    blank_screen(&mut screen, start_line);
    screen.move_cursor(start_line, 0);
    let choices = matching::compute_matches(&choices, search.as_slice());
    let mut i = 0;
    for choice in choices.iter() {
      screen.tty.writeln(choice.as_slice());
      if i >= visible_choices {
        break;
      }
      i += 1;
    }

    match screen.tty.getchar() {
      Char(x) => search.push(x),
      Backspace => { search.pop(); }
      Control('h') => { search.pop(); }
      Control('u') => { search.clear(); }
      Control('c') => { return; }
      Enter => {
        println!("{}", choices[0]);
        break;
      }
      _ => panic!("Unexpected input"),
    }
  }
}

fn blank_screen(screen: &mut screen::Screen, start_line: u16) {
  screen.move_cursor(start_line, 0);
  let mut i = 0;
  while i < screen.height {
    let mut j = 0;
    while j < screen.width {
      screen.tty.write(" ".as_bytes());
      j += 1;
    }
    i += 1;
  }
}

struct Args {
  initial_search: String,
  help: bool,
}

fn parse_args(args: &Vec<String>) -> Option<Args> {
  let opts = [
    getopts::optflag("h", "help", "Show this message"),
    getopts::optopt("s", "search", "Specify an initial search string", "SEARCH"),
  ];

  let matches = match getopts::getopts(args.tail(), &opts) {
    Ok(m) => m,
    Err(f) => {
      println!("{}", f);
      print_usage(args[0].as_slice(), &opts);
      return None;
    }
  };

  let initial_search = match matches.opt_str("search") {
    Some(x) => x.clone(),
    None => String::from_str(""),
  };

  let help = matches.opt_present("help");
  if help {
    print_usage(args[0].as_slice(), &opts);
  }

  Some(Args { initial_search: initial_search, help: help })
}

fn print_usage(program: &str, opts: &[getopts::OptGroup]) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", getopts::usage(brief.as_slice(), opts));
}

fn read_choices() -> Vec<String> {
  let mut stdin = io::stdio::stdin();
  let mut lines = Vec::new();

  loop {
    match stdin.read_line() {
      Ok(x) => lines.push(String::from_str(x.as_slice().trim())),
      Err(_) => break,
    }
  }

  lines
}
<commit_msg>Add crude indexing<commit_after>#![allow(unstable, dead_code)]

extern crate getopts;
extern crate libc;

mod ansi;
mod matching;
mod screen;

use std::os;
use std::io;
use std::cmp::min;
use screen::Screen;
use screen::Key::*;

fn main() {
  let input_args = os::args();

  let args = match parse_args(&input_args) {
    Some(args) => args,
    None => {
      os::set_exit_status(1);
      return;
    },
  };

  if args.help { return; }

  let mut screen = Screen::open_screen();
  let mut index = 0;

  let choices = read_choices();
  let visible_choices = min(20, screen.height - 1);

  let mut search: String = args.initial_search.clone();
  let start_line = screen.height - visible_choices - 1;
  loop {
    blank_screen(&mut screen, start_line);
    screen.move_cursor(start_line, 0);
    let choices = matching::compute_matches(&choices, search.as_slice());
    let mut i = 0;
    for choice in choices.iter() {
      screen.tty.writeln(choice.as_slice());
      if i >= visible_choices {
        break;
      }
      i += 1;
    }

    match screen.tty.getchar() {
      Char(x) => search.push(x),
      Backspace => { search.pop(); }
      Control('h') => { search.pop(); }
      Control('u') => { search.clear(); }
      Control('c') => { return; }
      Control('n') => { index += 1; }
      Control('p') => { index -= 1; }
      Enter => {
        println!("{}", choices[index]);
        break;
      }
      _ => panic!("Unexpected input"),
    }
  }
}

fn blank_screen(screen: &mut screen::Screen, start_line: u16) {
  screen.move_cursor(start_line, 0);
  let mut i = 0;
  while i < screen.height {
    let mut j = 0;
    while j < screen.width {
      screen.tty.write(" ".as_bytes());
      j += 1;
    }
    i += 1;
  }
}

struct Args {
  initial_search: String,
  help: bool,
}

fn parse_args(args: &Vec<String>) -> Option<Args> {
  let opts = [
    getopts::optflag("h", "help", "Show this message"),
    getopts::optopt("s", "search", "Specify an initial search string", "SEARCH"),
  ];

  let matches = match getopts::getopts(args.tail(), &opts) {
    Ok(m) => m,
    Err(f) => {
      println!("{}", f);
      print_usage(args[0].as_slice(), &opts);
      return None;
    }
  };

  let initial_search = match matches.opt_str("search") {
    Some(x) => x.clone(),
    None => String::from_str(""),
  };

  let help = matches.opt_present("help");
  if help {
    print_usage(args[0].as_slice(), &opts);
  }

  Some(Args { initial_search: initial_search, help: help })
}

fn print_usage(program: &str, opts: &[getopts::OptGroup]) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", getopts::usage(brief.as_slice(), opts));
}

fn read_choices() -> Vec<String> {
  let mut stdin = io::stdio::stdin();
  let mut lines = Vec::new();

  loop {
    match stdin.read_line() {
      Ok(x) => lines.push(String::from_str(x.as_slice().trim())),
      Err(_) => break,
    }
  }

  lines
}
<|endoftext|>


<commit_before>use std::fs;
use std::io;
use std::path::Path;

use rand::seq::SliceRandom;

const FORTUNES: &str = "fortunes";

fn main() {
    let fortunes = fortunes_from_file(FORTUNES).expect("Failed to read fortunes");

    let mut rng = rand::thread_rng();
    let picked_fortune = fortunes.choose(&mut rng).expect("No fortunes available");

    println!("{}", picked_fortune);
}

/// Read fortunes from a fortunes text file.
fn fortunes_from_file<P: AsRef<Path>>(path: P) -> io::Result<Vec<String>> {
    let data = fs::read_to_string(path)?;

    let mut fortunes = Vec::new();
    for fortune in data.split("\n%\n") {
        fortunes.push(fortune.to_string());
    }

    Ok(fortunes)
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_fortunes_from_file() {
        let fortunes = fortunes_from_file(FORTUNES).unwrap();

        assert_eq!(fortunes[15], "To teach is to learn.");
        assert_eq!(fortunes[39], "Write a wise saying and your name will live forever.\n                -- Anonymous");
        assert_eq!(fortunes[98], "All rights reserved.");
    }
}<commit_msg>Fortune: Implement TCP server<commit_after>use std::fs;
use std::io;
use std::io::Write;
use std::net::TcpListener;
use std::net::TcpStream;
use std::path::Path;

use rand::seq::SliceRandom;

const ADDR: &str = "127.0.0.1:5000";
const FORTUNES: &str = "fortunes";

fn main() {
    let fortunes = fortunes_from_file(FORTUNES).expect("Failed to read fortunes");

    eprintln!("Listening on {}", ADDR);
    let listener = TcpListener::bind(ADDR).expect("Failed to bind address");

    for incomming in listener.incoming() {
        let stream = match incomming {
            Err(e) => {
                eprintln!("Accept failed: {}", e);
                continue;
            },
            Ok(stream) => stream,
        };

        let peer_addr = stream.peer_addr().expect("Failed to get peer address");
        eprintln!("Accepted connection from {}", peer_addr);
        handle_stream(stream, &fortunes);
    }
}

/// Read fortunes from a fortunes text file.
fn fortunes_from_file<P: AsRef<Path>>(path: P) -> io::Result<Vec<String>> {
    let data = fs::read_to_string(path)?;

    let mut fortunes = Vec::new();
    for fortune in data.split("\n%\n") {
        fortunes.push(fortune.to_string());
    }

    Ok(fortunes)
}

/// Handle a single connection stream.
fn handle_stream(mut stream: TcpStream, fortunes: &[String]) {
    let mut rng = rand::thread_rng();
    let picked_fortune = fortunes.choose(&mut rng).expect("No fortunes available");

    writeln!(stream, "{}", picked_fortune).expect("Write failed");
}

#[cfg(test)]
mod test {
    use std::io::Read;

    use super::*;

    #[test]
    fn test_fortunes_from_file() {
        let fortunes = fortunes_from_file(FORTUNES).unwrap();

        assert_eq!(fortunes[15], "To teach is to learn.");
        assert_eq!(fortunes[39], "Write a wise saying and your name will live forever.\n                -- Anonymous");
        assert_eq!(fortunes[98], "All rights reserved.");
    }

    #[test]
    fn test_handle_steam() {
        let fortunes = vec![String::from("Reply hazy, try again.")];

        let listener = TcpListener::bind("127.0.0.1:0").expect("Failed to bind address");
        let local_addr = listener.local_addr().expect("Failed to get local address");

        std::thread::spawn(move || {
            let (stream, _) = listener.accept().expect("Failed to accept connection");

            handle_stream(stream, &fortunes);
        });

        let mut client = TcpStream::connect(local_addr).unwrap();
        let mut buf = String::new();
        client.read_to_string(&mut buf).expect("Read failed");

        assert_eq!(buf, "Reply hazy, try again.\n");
    }
}<|endoftext|>


<commit_before>use std::io;
use std::io::BufReader;
use std::io::BufRead;
use std::collections::HashMap;

mod request;
mod turn;
mod region;
mod superregion;

use request::Request;
use turn::Turn;
use region::Region;
use superregion::SuperRegion;

fn main()
{
	std::env::set_var("RUST_BACKTRACE","1");
	let input = io::stdin();
	let input = BufReader::new(input);
	let input = input.lines();
	let input = input
		.map(Result::unwrap)
		.filter(|s|s.len()>0)
		.map(|s|s.parse::<Request>())
		.map(Result::unwrap);

	let mut name_you = String::new();
	let mut name_other = String::new();

	let mut sup_regions: HashMap<usize,SuperRegion> = HashMap::new();
	let mut regions: HashMap<usize,Region> = HashMap::new();

	let mut armies_left = 0;

	for req in input
	{
		use request::Request::*;
		match req
		{
			ListSuperRegions(sregions) =>
			{
				sup_regions.extend(sregions.iter().map(|sr|(sr.id,sr.clone())));
			},
			ListRegions(regs) =>
			{
				regions.extend(regs.iter().map(|r|(r.id,r.clone())));
			},
			ListNeighbours(relations) =>
			{
				for (id,neighbours) in relations
				{
					regions.get_mut(&id).unwrap().neighbours.extend(neighbours.iter());
					for n in neighbours
					{
						regions.get_mut(&n).unwrap().neighbours.push(id);
					}
				}
			},
			RequestStartingRegions(_avail) =>
			{
				println!("give me randomly");
			},
			SettingNameYou(name) => name_you = name,
			SettingNameOther(name) => name_other = name,
			SettingArmies(count) =>
			{
				armies_left = count;
			},
			UpdateMap(updates) =>
			{
				for u in updates
				{
					let mut reg = regions.get_mut(&u.0).unwrap();
					reg.player = u.1;
					reg.count = u.2;
				}
			},
			TurnOther(_turns) =>
			{
				//TODO: what the hell?
			},
			TurnPlace =>
			{
				let mut v = Vec::new();
				{
					let regs = regions.values()
						.filter(|r|r.player == name_you)
						.filter(|r|r.neighbours.iter()
							.map(|o|regions.get(o))
							.map(Option::unwrap)
							.filter(|o|o.player == name_other)
							.any(|o|o.count > r.count)
						).collect::<Vec<_>>();

					if armies_left > regs.len()
					{
						let count = if regs.len() > 0
						{
							armies_left/regs.len()
						}
						else
						{
							1
						};
						for r in regs.iter()
						{
							v.push(
								Turn::Place
								{
									name: name_you.clone(),
									region: r.id,
									count: count,
								}
							);
							armies_left -= count;
						}
						for r in regions.values()
							.filter(|r|r.player == name_you)
						{
							if armies_left <= 0
							{
								break;
							}
							v.push(
								Turn::Place
								{
									name: name_you.clone(),
									region: r.id,
									count: 1,
								}
							);
							armies_left -= 1;
						}
					}
					if v.len() <= 0
					{
						v.push(Turn::Noop);
					}
					let v = v.iter()
						.filter(|t|
							if let &&Turn::Place{count,..} = t
							{
								count > 0
							}
							else
							{
								true
							}
						)
						.map(|t|format!("{}",t))
						.collect::<Vec<_>>();
					println!("{}",v.join(","));
				}
				for r in v.iter()
				{
					if let &Turn::Place{region,count,..} = r
					{
						regions.get_mut(&region).unwrap().count += count;
					}
				}
			},
			TurnArmies =>
			{
				let mut v = Vec::new();
				let regs = regions.values()
					.filter(|r|r.player == name_you)
					.filter_map(|r|
					{
						let x =r.neighbours.iter()
						.map(|o|regions.get(o))
						.map(Option::unwrap)
						.filter(|o|o.player != name_you)
						.filter(|o|o.count*5 <= r.count*2)
						.next()
						.map(|o|o.id);
						if let Some(x) = x
						{
							Some(((r.id,(r.count*2)/3),x))
						}
						else
						{
							None
						}
					})
					.collect::<Vec<_>>();
				for ((r,count),n) in regs
				{
					v.push(
						Turn::Turn
						{
							name: name_you.clone(),
							source: r,
							target: n,
							count: count,
						}
					);
				}
				if v.len() <= 0
				{
					v.push(Turn::Noop);
				}
				let v = v.iter().map(|t|format!("{}",t)).collect::<Vec<_>>();
				println!("{}",v.join(","));
			},
		}
	}
}

<commit_msg>small number adjustment<commit_after>use std::io;
use std::io::BufReader;
use std::io::BufRead;
use std::collections::HashMap;

mod request;
mod turn;
mod region;
mod superregion;

use request::Request;
use turn::Turn;
use region::Region;
use superregion::SuperRegion;

fn main()
{
	std::env::set_var("RUST_BACKTRACE","1");
	let input = io::stdin();
	let input = BufReader::new(input);
	let input = input.lines();
	let input = input
		.map(Result::unwrap)
		.filter(|s|s.len()>0)
		.map(|s|s.parse::<Request>())
		.map(Result::unwrap);

	let mut name_you = String::new();
	let mut name_other = String::new();

	let mut sup_regions: HashMap<usize,SuperRegion> = HashMap::new();
	let mut regions: HashMap<usize,Region> = HashMap::new();

	let mut armies_left = 0;

	for req in input
	{
		use request::Request::*;
		match req
		{
			ListSuperRegions(sregions) =>
			{
				sup_regions.extend(sregions.iter().map(|sr|(sr.id,sr.clone())));
			},
			ListRegions(regs) =>
			{
				regions.extend(regs.iter().map(|r|(r.id,r.clone())));
			},
			ListNeighbours(relations) =>
			{
				for (id,neighbours) in relations
				{
					regions.get_mut(&id).unwrap().neighbours.extend(neighbours.iter());
					for n in neighbours
					{
						regions.get_mut(&n).unwrap().neighbours.push(id);
					}
				}
			},
			RequestStartingRegions(_avail) =>
			{
				println!("give me randomly");
			},
			SettingNameYou(name) => name_you = name,
			SettingNameOther(name) => name_other = name,
			SettingArmies(count) =>
			{
				armies_left = count;
			},
			UpdateMap(updates) =>
			{
				for u in updates
				{
					let mut reg = regions.get_mut(&u.0).unwrap();
					reg.player = u.1;
					reg.count = u.2;
				}
			},
			TurnOther(_turns) =>
			{
				//TODO: what the hell?
			},
			TurnPlace =>
			{
				let mut v = Vec::new();
				{
					let regs = regions.values()
						.filter(|r|r.player == name_you)
						.filter(|r|r.neighbours.iter()
							.map(|o|regions.get(o))
							.map(Option::unwrap)
							.filter(|o|o.player == name_other)
							.any(|o|o.count >= r.count)
						).collect::<Vec<_>>();

					if armies_left > regs.len()
					{
						let count = if regs.len() > 0
						{
							armies_left/regs.len()
						}
						else
						{
							1
						};
						for r in regs.iter()
						{
							v.push(
								Turn::Place
								{
									name: name_you.clone(),
									region: r.id,
									count: count,
								}
							);
							armies_left -= count;
						}
						for r in regions.values()
							.filter(|r|r.player == name_you)
						{
							if armies_left <= 0
							{
								break;
							}
							v.push(
								Turn::Place
								{
									name: name_you.clone(),
									region: r.id,
									count: 1,
								}
							);
							armies_left -= 1;
						}
					}
					if v.len() <= 0
					{
						v.push(Turn::Noop);
					}
					let v = v.iter()
						.filter(|t|
							if let &&Turn::Place{count,..} = t
							{
								count > 0
							}
							else
							{
								true
							}
						)
						.map(|t|format!("{}",t))
						.collect::<Vec<_>>();
					println!("{}",v.join(","));
				}
				for r in v.iter()
				{
					if let &Turn::Place{region,count,..} = r
					{
						regions.get_mut(&region).unwrap().count += count;
					}
				}
			},
			TurnArmies =>
			{
				let mut v = Vec::new();
				let regs = regions.values()
					.filter(|r|r.player == name_you)
					.filter_map(|r|
					{
						let x =r.neighbours.iter()
						.map(|o|regions.get(o))
						.map(Option::unwrap)
						.filter(|o|o.player != name_you)
						.filter(|o|o.count*5 <= r.count*2)
						.next()
						.map(|o|o.id);
						if let Some(x) = x
						{
							Some(((r.id,(r.count*2)/3),x))
						}
						else
						{
							None
						}
					})
					.collect::<Vec<_>>();
				for ((r,count),n) in regs
				{
					v.push(
						Turn::Turn
						{
							name: name_you.clone(),
							source: r,
							target: n,
							count: count,
						}
					);
				}
				if v.len() <= 0
				{
					v.push(Turn::Noop);
				}
				let v = v.iter().map(|t|format!("{}",t)).collect::<Vec<_>>();
				println!("{}",v.join(","));
			},
		}
	}
}

<|endoftext|>


<commit_before>#![feature(libc)]

use std::collections::BTreeMap;

extern crate numbat;

mod forza;
mod plugins;

use forza::ForzaPlugin;

fn start_plugin<T: ForzaPlugin>(mut plugin: T) {
  plugin.start();
}

fn main() {
  let mut emitter = numbat::Emitter::new(BTreeMap::new(), "forza");
  emitter.connect("tcp://127.0.0.1:1337");

  let heartbeat = plugins::load_average::LoadAverage::new(emitter);
  start_plugin(heartbeat);
}
<commit_msg>Start both plugins in threads, clone emitters<commit_after>extern crate libc;
extern crate numbat;

mod forza;
mod plugins;

use std::thread;
use std::collections::BTreeMap;
use forza::ForzaPlugin;

fn start_plugin<T: ForzaPlugin>(mut plugin: T) {
  plugin.start();
}

fn main() {
  let mut emitter = numbat::Emitter::new(BTreeMap::new(), "forza");
  emitter.connect("tcp://127.0.0.1:1337");

  let heartbeat = plugins::heartbeat::Heartbeat::new(emitter.clone());
  thread::spawn(|| {
    start_plugin(heartbeat);
  });

  let load_average = plugins::load_average::LoadAverage::new(emitter.clone());
  thread::spawn(|| {
    start_plugin(load_average);
  }).join();
}
<|endoftext|>


<commit_before>use clap::Clap;
use std::cmp::Ordering;
use std::fs::File;

fn normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)$").unwrap();
    }
    let mut col = RE
        .replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace("(", "")
        .replace(")", "")
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
        .replace(",", "_")
        .replace("&", "_")
        .replace("#", "");
    if !col.chars().next().map(char::is_alphabetic).unwrap_or(true) {
        col = format!("c_{}", col)
    }
    col
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: &str,
    delimiter: u8,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::ReaderBuilder::new()
        .flexible(true)
        .delimiter(delimiter)
        .from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).expect("tx.prepare() failed");
        while let Some(row) = records.next() {
            let mut row = row?;
            match row.len().cmp(&normalized_cols.len()) {
                Ordering::Less => {
                    for _ in 0..normalized_cols.len() - row.len() {
                        row.push_field("");
                    }
                }
                Ordering::Greater => {
                    panic!("Too many fields on row {}, fields: {:?}", num_rows + 1, row);
                }
                Ordering::Equal => {}
            }
            stmt.execute(rusqlite::params_from_iter(&row)).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({}) from {:?}",
        num_rows,
        table_name,
        normalized_cols.join(", "),
        path,
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;

    let mut table = comfy_table::Table::new();
    table.load_preset("││──╞═╪╡┆    ┬┴┌┐└┘");
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
    let mut title_row = comfy_table::Row::new();
    for col in stmt.column_names() {
        title_row.add_cell(comfy_table::Cell::new(col));
    }
    table.set_header(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = comfy_table::Row::new();
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(comfy_table::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    println!("{}", table);
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else if line.starts_with(".schema") {
        _handle_query(
            conn,
            "SELECT sql AS schema FROM sqlite_master WHERE name like 't%'",
        )
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

fn install_udfs(c: &mut rusqlite::Connection) -> Result<(), Box<dyn std::error::Error>> {
    c.create_scalar_function(
        "regexp_extract",
        3,
        rusqlite::functions::FunctionFlags::SQLITE_UTF8
            | rusqlite::functions::FunctionFlags::SQLITE_DETERMINISTIC,
        |ctx| {
            let re = ctx.get_or_create_aux(
                0,
                |vr| -> Result<_, Box<dyn std::error::Error + Send + Sync + 'static>> {
                    Ok(regex::Regex::new(vr.as_str()?)?)
                },
            )?;
            let value = ctx.get::<Box<str>>(1)?;
            let replacement = ctx.get::<Box<str>>(2)?;

            let caps = match re.captures(&value) {
                Some(caps) => caps,
                None => return Ok("".to_string()),
            };
            let mut dest = String::new();
            caps.expand(&replacement, &mut dest);
            Ok(dest)
        },
    )?;

    Ok(())
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 5] = [b' ', b'(', b')', b',', b'.'];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    type Hint = String;

    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::validate::Validator for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

#[derive(Clap)]
struct Opts {
    #[clap(long, about = "Use ',' as the delimiter for the CSV")]
    comma: bool,
    #[clap(long, about = "Use '|' as the delimiter for the CSV")]
    pipe: bool,
    #[clap(long, about = "Use '\\t' as the delimiter for the CSV")]
    tab: bool,
    #[clap(long, about = "Use ';' as the delimiter for the CSV")]
    semicolon: bool,

    #[clap()]
    paths: Vec<String>,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let opts = Opts::parse();

    let delim = match (opts.comma, opts.pipe, opts.tab, opts.semicolon) {
        (true, false, false, false) | (false, false, false, false) => b',',
        (false, true, false, false) => b'|',
        (false, false, true, false) => b'\t',
        (false, false, false, true) => b';',
        _ => {
            eprintln!("Can't pass more than one of --comma, --pipe, and --tab");
            std::process::exit(1);
        }
    };

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct",
        "select",
        "from",
        "group",
        "by",
        "order",
        "where",
        "count",
        "limit",
        "offset",
        "length",
        "coalesce",
        "regexp_extract",
        ".export",
        ".schema",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if opts.paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", &opts.paths[0], delim)?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in opts.paths.iter().enumerate() {
            let mut col_names =
                _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path, delim)?;
            base_words.append(&mut col_names);
        }
    }

    install_udfs(&mut conn)?;

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }

    Ok(())
}

#[cfg(test)]
mod test {
    use super::normalize_col;

    #[test]
    fn test_normalize_col() {
        for (value, expected) in &[
            ("", ""),
            ("abc", "abc"),
            ("(S)AO", "sao"),
            ("abc (123)", "abc"),
            ("2/6/2000", "c_2_6_2000"),
            ("COMBO#", "combo"),
        ] {
            assert_eq!(&&normalize_col(value), expected);
        }
    }
}
<commit_msg>Allow non-utf8 data to be loaded (#593)<commit_after>use clap::Clap;
use std::cmp::Ordering;
use std::fs::File;

fn normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)$").unwrap();
    }
    let mut col = RE
        .replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace("(", "")
        .replace(")", "")
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
        .replace(",", "_")
        .replace("&", "_")
        .replace("#", "");
    if !col.chars().next().map(char::is_alphabetic).unwrap_or(true) {
        col = format!("c_{}", col)
    }
    col
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: &str,
    delimiter: u8,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::ReaderBuilder::new()
        .flexible(true)
        .delimiter(delimiter)
        .from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.byte_records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).expect("tx.prepare() failed");
        while let Some(row) = records.next() {
            let mut row = row?;
            match row.len().cmp(&normalized_cols.len()) {
                Ordering::Less => {
                    for _ in 0..normalized_cols.len() - row.len() {
                        row.push_field(b"");
                    }
                }
                Ordering::Greater => {
                    panic!("Too many fields on row {}, fields: {:?}", num_rows + 1, row);
                }
                Ordering::Equal => {}
            }
            stmt.execute(rusqlite::params_from_iter(
                row.iter().map(String::from_utf8_lossy),
            ))
            .unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({}) from {:?}",
        num_rows,
        table_name,
        normalized_cols.join(", "),
        path,
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;

    let mut table = comfy_table::Table::new();
    table.load_preset("││──╞═╪╡┆    ┬┴┌┐└┘");
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
    let mut title_row = comfy_table::Row::new();
    for col in stmt.column_names() {
        title_row.add_cell(comfy_table::Cell::new(col));
    }
    table.set_header(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = comfy_table::Row::new();
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(comfy_table::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    println!("{}", table);
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else if line.starts_with(".schema") {
        _handle_query(
            conn,
            "SELECT sql AS schema FROM sqlite_master WHERE name like 't%'",
        )
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

fn install_udfs(c: &mut rusqlite::Connection) -> Result<(), Box<dyn std::error::Error>> {
    c.create_scalar_function(
        "regexp_extract",
        3,
        rusqlite::functions::FunctionFlags::SQLITE_UTF8
            | rusqlite::functions::FunctionFlags::SQLITE_DETERMINISTIC,
        |ctx| {
            let re = ctx.get_or_create_aux(
                0,
                |vr| -> Result<_, Box<dyn std::error::Error + Send + Sync + 'static>> {
                    Ok(regex::Regex::new(vr.as_str()?)?)
                },
            )?;
            let value = ctx.get::<Box<str>>(1)?;
            let replacement = ctx.get::<Box<str>>(2)?;

            let caps = match re.captures(&value) {
                Some(caps) => caps,
                None => return Ok("".to_string()),
            };
            let mut dest = String::new();
            caps.expand(&replacement, &mut dest);
            Ok(dest)
        },
    )?;

    Ok(())
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 5] = [b' ', b'(', b')', b',', b'.'];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    type Hint = String;

    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::validate::Validator for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

#[derive(Clap)]
struct Opts {
    #[clap(long, about = "Use ',' as the delimiter for the CSV")]
    comma: bool,
    #[clap(long, about = "Use '|' as the delimiter for the CSV")]
    pipe: bool,
    #[clap(long, about = "Use '\\t' as the delimiter for the CSV")]
    tab: bool,
    #[clap(long, about = "Use ';' as the delimiter for the CSV")]
    semicolon: bool,

    #[clap()]
    paths: Vec<String>,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let opts = Opts::parse();

    let delim = match (opts.comma, opts.pipe, opts.tab, opts.semicolon) {
        (true, false, false, false) | (false, false, false, false) => b',',
        (false, true, false, false) => b'|',
        (false, false, true, false) => b'\t',
        (false, false, false, true) => b';',
        _ => {
            eprintln!("Can't pass more than one of --comma, --pipe, and --tab");
            std::process::exit(1);
        }
    };

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct",
        "select",
        "from",
        "group",
        "by",
        "order",
        "where",
        "count",
        "limit",
        "offset",
        "length",
        "coalesce",
        "regexp_extract",
        ".export",
        ".schema",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if opts.paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", &opts.paths[0], delim)?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in opts.paths.iter().enumerate() {
            let mut col_names =
                _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path, delim)?;
            base_words.append(&mut col_names);
        }
    }

    install_udfs(&mut conn)?;

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }

    Ok(())
}

#[cfg(test)]
mod test {
    use super::normalize_col;

    #[test]
    fn test_normalize_col() {
        for (value, expected) in &[
            ("", ""),
            ("abc", "abc"),
            ("(S)AO", "sao"),
            ("abc (123)", "abc"),
            ("2/6/2000", "c_2_6_2000"),
            ("COMBO#", "combo"),
        ] {
            assert_eq!(&&normalize_col(value), expected);
        }
    }
}
<|endoftext|>


<commit_before>//!                         The Unlicense
//! This is free and unencumbered software released into the public domain.
//!
//! Anyone is free to copy, modify, publish, use, compile, sell, or
//! distribute this software, either in source code form or as a compiled
//! binary, for any purpose, commercial or non-commercial, and by any
//! means.
//!
//! In jurisdictions that recognize copyright laws, the author or authors
//! of this software dedicate any and all copyright interest in the
//! software to the public domain. We make this dedication for the benefit
//! of the public at large and to the detriment of our heirs and
//! successors. We intend this dedication to be an overt act of
//! relinquishment in perpetuity of all present and future rights to this
//! software under copyright law.
//!
//! THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
//! EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
//! MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
//! IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
//! OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
//! ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
//! OTHER DEALINGS IN THE SOFTWARE.
//!
//! For more information, please refer to <http://unlicense.org>
//!

extern crate iron;
extern crate router;
extern crate hyper;
#[macro_use]
extern crate lazy_static;
extern crate regex;
extern crate hyper_rustls;
extern crate urlencoded;

use std::vec::Vec;
use hyper::Client as HyperClient;
use hyper::net::HttpsConnector;
use router::{Router, NoRoute};
use iron::prelude::*;
use iron::status;
use iron::AfterMiddleware;
use regex::Regex;
use std::io::Read;
use hyper_rustls::TlsClient as RustlsClient;
use urlencoded::{UrlEncodedQuery, UrlDecodingError};
use std::collections::HashMap;
use std::option::Option;
use iron::headers::ContentType;
use std::ops::Deref;

fn verify_md5(md5: &str) -> bool {
    lazy_static! {
        static ref MD5_RE: Regex = Regex::new("^[0-9a-f]{32}$").unwrap();
    }

    MD5_RE.is_match(md5)
}

struct AvatarError {
    pub status: Option<status::Status>,
}

struct Avatar {
    pub email_md5: String,
    pub size: Option<i64>,
    pub content_type: ContentType,
    pub buf: Vec<u8>,
}

type AvatarResult = Result<Avatar, AvatarError>;

impl Avatar {
    fn make_origin_url(email_md5: &str, size: Option<i64>, default: Option<String>) -> String {
        let mut origin_url = format!("https://secure.gravatar.com/avatar/{}", email_md5);

        let mut params: Vec<String> = Vec::new();

        if let Some(m) = size {
            params.push(format!("s={}", m));
        }

        if let Some(m) = default {
            params.push(format!("d={}", m));
        }

        if params.len() > 0 {
            origin_url.push_str("?");
            origin_url.push_str(params.join("&").as_str());
        }

        origin_url
    }

    pub fn fetch(email_md5: &str, size: Option<i64>, default: Option<String>) -> AvatarResult {

        lazy_static! {
            static ref HYPER_CLIENT: HyperClient = HyperClient::with_connector(HttpsConnector::new(RustlsClient::new()));
        }

        let origin_url = Avatar::make_origin_url(email_md5, size.clone(), default.clone());

        let mut origin_res = try!(HYPER_CLIENT.get(origin_url.as_str())
            .send()
            .map_err(|_| AvatarError { status: None }));

        if origin_res.status != hyper::Ok {
            return Err(AvatarError { status: Some(origin_res.status) });
        }

        let mut avatar_body: Vec<u8> = Vec::new();

        try!(origin_res.read_to_end(&mut avatar_body)
            .map_err(|_| AvatarError { status: Some(status::BadGateway) }));

        let content_type = try!(origin_res.headers
                .get::<ContentType>()
                .ok_or(AvatarError { status: Some(status::BadGateway) }))
            .clone();

        Ok(Avatar {
            email_md5: email_md5.to_string(),
            size: size.clone(),
            content_type: content_type,
            buf: avatar_body,
        })
    }
}

fn handler(req: &mut Request) -> IronResult<Response> {
    let email_md5 = req.extensions
        .get::<Router>()
        .unwrap()
        .find("email_md5")
        .unwrap_or("")
        .to_lowercase();

    if !verify_md5(email_md5.as_str()) {
        return Ok(Response::with((status::NotFound)));
    }

    let maybe_query = req.get::<UrlEncodedQuery>()
        .map(|val| Some(val))
        .unwrap_or_else(|err: UrlDecodingError| {
            match err {
                UrlDecodingError::EmptyQuery => Some(HashMap::new()),
                _ => None,
            }
        });

    let query = match maybe_query {
        Some(m) => m,
        None => return Ok(Response::with((status::BadRequest))),
    };

    let size_str = query.get("s")
        .map(|s_vec| s_vec.clone())
        .map(|s_vec| { if s_vec.len() > 0 { Some(s_vec) } else { None } })
        .unwrap_or(None)
        .map(|s_vec| s_vec[0].clone());

    let default = query.get("d")
        .map(|d_vec| d_vec.clone())
        .map(|d_vec| { if d_vec.len() > 0 { Some(d_vec) } else { None } })
        .unwrap_or(None)
        .map(|d_vec| d_vec[0].clone());

    let size: Option<i64> = match size_str {
        Some(m) => {
            match m.parse::<i64>() {
                Ok(m) => Some(m),
                Err(_) => return Ok(Response::with((status::BadRequest))),
            }
        }
        None => None,
    };

    let avatar = match Avatar::fetch(email_md5.as_str(), size, default) {
        Ok(v) => v,
        Err(_) => return Ok(Response::with((status::BadGateway))),
    };

    Ok(Response::with((avatar.content_type.deref().clone(), status::Ok, avatar.buf.clone())))
}

struct CustomErrorMsg;

impl CustomErrorMsg {
    fn alter_response(&self, res: Response) -> Response {
        if res.status == Some(status::NotFound) {
            Response::with((status::NotFound, "HTTP 404: Not Found."))

        } else if res.status == Some(status::BadRequest) {
            Response::with((status::BadRequest, "HTTP 400: Bad Request."))

        } else if res.status == Some(status::BadGateway) {
            Response::with((status::BadGateway, "HTTP 502: Bad Gateway."))

        } else {
            res

        }
    }
}

impl AfterMiddleware for CustomErrorMsg {
    fn after(&self, _: &mut Request, res: Response) -> IronResult<Response> {
        Ok(self.alter_response(res))
    }

    fn catch(&self, _: &mut Request, err: IronError) -> IronResult<Response> {
        let res = err.response;
        let res = self.alter_response(res);

        if let Some(_) = err.error.downcast::<NoRoute>() {
            return Ok(res);
        }

        Err(IronError {
            error: err.error,
            response: res,
        })
    }
}

struct XPoweredBy;

impl XPoweredBy {
    fn add_header(&self, res: &mut Response) {
        res.headers.set_raw("X-Powered-By", vec![b"Rust".to_vec()]);
    }
}

impl AfterMiddleware for XPoweredBy {
    fn after(&self, _: &mut Request, mut res: Response) -> IronResult<Response> {
        self.add_header(&mut res);

        Ok(res)
    }

    fn catch(&self, _: &mut Request, err: IronError) -> IronResult<Response> {
        let mut res = err.response;
        let inner_err = err.error;
        self.add_header(&mut res);

        Err(IronError {
            error: inner_err,
            response: res,
        })
    }
}

fn main() {
    let mut router = Router::new();

    router.get("/avatar/:email_md5", handler, "email_md5");

    let mut chain = Chain::new(router);
    chain.link_after(CustomErrorMsg);
    chain.link_after(XPoweredBy);

    Iron::new(chain).http("localhost:3000").unwrap();
}
<commit_msg>Implement Debug, Display, and Error for AvatarError.<commit_after>//!                         The Unlicense
//! This is free and unencumbered software released into the public domain.
//!
//! Anyone is free to copy, modify, publish, use, compile, sell, or
//! distribute this software, either in source code form or as a compiled
//! binary, for any purpose, commercial or non-commercial, and by any
//! means.
//!
//! In jurisdictions that recognize copyright laws, the author or authors
//! of this software dedicate any and all copyright interest in the
//! software to the public domain. We make this dedication for the benefit
//! of the public at large and to the detriment of our heirs and
//! successors. We intend this dedication to be an overt act of
//! relinquishment in perpetuity of all present and future rights to this
//! software under copyright law.
//!
//! THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
//! EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
//! MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
//! IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
//! OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
//! ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
//! OTHER DEALINGS IN THE SOFTWARE.
//!
//! For more information, please refer to <http://unlicense.org>
//!

extern crate iron;
extern crate router;
extern crate hyper;
#[macro_use]
extern crate lazy_static;
extern crate regex;
extern crate hyper_rustls;
extern crate urlencoded;

use std::vec::Vec;
use hyper::Client as HyperClient;
use hyper::net::HttpsConnector;
use router::{Router, NoRoute};
use iron::prelude::*;
use iron::status;
use iron::AfterMiddleware;
use regex::Regex;
use std::io::Read;
use hyper_rustls::TlsClient as RustlsClient;
use urlencoded::{UrlEncodedQuery, UrlDecodingError};
use std::collections::HashMap;
use std::option::Option;
use iron::headers::ContentType;
use std::ops::Deref;
use std::fmt;
use std::error::Error;

fn verify_md5(md5: &str) -> bool {
    lazy_static! {
        static ref MD5_RE: Regex = Regex::new("^[0-9a-f]{32}$").unwrap();
    }

    MD5_RE.is_match(md5)
}

#[derive(Debug)]
struct AvatarError {
    pub status: Option<status::Status>,
}

impl fmt::Display for AvatarError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Cannot Fetch Avatar!")
    }
}

impl Error for AvatarError {
    fn description(&self) -> &str {
        "Cannot Fetch Avatar!"
    }
}

struct Avatar {
    pub email_md5: String,
    pub size: Option<i64>,
    pub content_type: ContentType,
    pub buf: Vec<u8>,
}

type AvatarResult = Result<Avatar, AvatarError>;

impl Avatar {
    fn make_origin_url(email_md5: &str, size: Option<i64>, default: Option<String>) -> String {
        let mut origin_url = format!("https://secure.gravatar.com/avatar/{}", email_md5);

        let mut params: Vec<String> = Vec::new();

        if let Some(m) = size {
            params.push(format!("s={}", m));
        }

        if let Some(m) = default {
            params.push(format!("d={}", m));
        }

        if params.len() > 0 {
            origin_url.push_str("?");
            origin_url.push_str(params.join("&").as_str());
        }

        origin_url
    }

    pub fn fetch(email_md5: &str, size: Option<i64>, default: Option<String>) -> AvatarResult {

        lazy_static! {
            static ref HYPER_CLIENT: HyperClient = HyperClient::with_connector(HttpsConnector::new(RustlsClient::new()));
        }

        let origin_url = Avatar::make_origin_url(email_md5, size.clone(), default.clone());

        let mut origin_res = try!(HYPER_CLIENT.get(origin_url.as_str())
            .send()
            .map_err(|_| AvatarError { status: None }));

        if origin_res.status != hyper::Ok {
            return Err(AvatarError { status: Some(origin_res.status) });
        }

        let mut avatar_body: Vec<u8> = Vec::new();

        try!(origin_res.read_to_end(&mut avatar_body)
            .map_err(|_| AvatarError { status: Some(status::BadGateway) }));

        let content_type = try!(origin_res.headers
                .get::<ContentType>()
                .ok_or(AvatarError { status: Some(status::BadGateway) }))
            .clone();

        Ok(Avatar {
            email_md5: email_md5.to_string(),
            size: size.clone(),
            content_type: content_type,
            buf: avatar_body,
        })
    }
}

fn handler(req: &mut Request) -> IronResult<Response> {
    let email_md5 = req.extensions
        .get::<Router>()
        .unwrap()
        .find("email_md5")
        .unwrap_or("")
        .to_lowercase();

    if !verify_md5(email_md5.as_str()) {
        return Ok(Response::with((status::NotFound)));
    }

    let maybe_query = req.get::<UrlEncodedQuery>()
        .map(|val| Some(val))
        .unwrap_or_else(|err: UrlDecodingError| {
            match err {
                UrlDecodingError::EmptyQuery => Some(HashMap::new()),
                _ => None,
            }
        });

    let query = match maybe_query {
        Some(m) => m,
        None => return Ok(Response::with((status::BadRequest))),
    };

    let size_str = query.get("s")
        .map(|s_vec| s_vec.clone())
        .map(|s_vec| { if s_vec.len() > 0 { Some(s_vec) } else { None } })
        .unwrap_or(None)
        .map(|s_vec| s_vec[0].clone());

    let default = query.get("d")
        .map(|d_vec| d_vec.clone())
        .map(|d_vec| { if d_vec.len() > 0 { Some(d_vec) } else { None } })
        .unwrap_or(None)
        .map(|d_vec| d_vec[0].clone());

    let size: Option<i64> = match size_str {
        Some(m) => {
            match m.parse::<i64>() {
                Ok(m) => Some(m),
                Err(_) => return Ok(Response::with((status::BadRequest))),
            }
        }
        None => None,
    };

    let avatar = match Avatar::fetch(email_md5.as_str(), size, default) {
        Ok(v) => v,
        Err(_) => return Ok(Response::with((status::BadGateway))),
    };

    Ok(Response::with((avatar.content_type.deref().clone(), status::Ok, avatar.buf.clone())))
}

struct CustomErrorMsg;

impl CustomErrorMsg {
    fn alter_response(&self, res: Response) -> Response {
        if res.status == Some(status::NotFound) {
            Response::with((status::NotFound, "HTTP 404: Not Found."))

        } else if res.status == Some(status::BadRequest) {
            Response::with((status::BadRequest, "HTTP 400: Bad Request."))

        } else if res.status == Some(status::BadGateway) {
            Response::with((status::BadGateway, "HTTP 502: Bad Gateway."))

        } else {
            res

        }
    }
}

impl AfterMiddleware for CustomErrorMsg {
    fn after(&self, _: &mut Request, res: Response) -> IronResult<Response> {
        Ok(self.alter_response(res))
    }

    fn catch(&self, _: &mut Request, err: IronError) -> IronResult<Response> {
        let res = err.response;
        let res = self.alter_response(res);

        if let Some(_) = err.error.downcast::<NoRoute>() {
            return Ok(res);
        }

        Err(IronError {
            error: err.error,
            response: res,
        })
    }
}

struct XPoweredBy;

impl XPoweredBy {
    fn add_header(&self, res: &mut Response) {
        res.headers.set_raw("X-Powered-By", vec![b"Rust".to_vec()]);
    }
}

impl AfterMiddleware for XPoweredBy {
    fn after(&self, _: &mut Request, mut res: Response) -> IronResult<Response> {
        self.add_header(&mut res);

        Ok(res)
    }

    fn catch(&self, _: &mut Request, err: IronError) -> IronResult<Response> {
        let mut res = err.response;
        let inner_err = err.error;
        self.add_header(&mut res);

        Err(IronError {
            error: inner_err,
            response: res,
        })
    }
}

fn main() {
    let mut router = Router::new();

    router.get("/avatar/:email_md5", handler, "email_md5");

    let mut chain = Chain::new(router);
    chain.link_after(CustomErrorMsg);
    chain.link_after(XPoweredBy);

    Iron::new(chain).http("localhost:3000").unwrap();
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
extern crate breezyvps;
extern crate simplelog;

use simplelog::{Config, TermLogger, WriteLogger, CombinedLogger, LogLevelFilter};
use std::fs::File;
use breezyvps::command;

fn sc_doctl(doctl_matches: &clap::ArgMatches) {
    if let Some(create_droplet_matches) = doctl_matches.subcommand_matches("create_droplet") {
        if let Some(name) = create_droplet_matches.value_of("name") {
            breezyvps::digitalocean::create_droplet_by_name(name);
        } else {
            println!("Missing required name parameter!");
        }
        return;
    }
    if let Some(destroy_droplet_matches) = doctl_matches.subcommand_matches("destroy_droplet") {
        if let Some(name) = destroy_droplet_matches.value_of("name") {
            breezyvps::digitalocean::destroy_droplet_by_name(name);
        } else {
            println!("Missing required name parameter!");
        }
        return;
    }
    if let Some(create_ssh_key_matches) = doctl_matches.subcommand_matches("create_sshkey") {
        if let Some(name) = create_ssh_key_matches.value_of("name") {
            breezyvps::digitalocean::create_sshkey(name);
        } else {
            println!("Missing required sshkey name parameter!");
        }
        return;
    }
}

fn sc_configure(configure_matches: &clap::ArgMatches) {
    if let Some(nginx_matches) = configure_matches.subcommand_matches("nginx") {
        if let Some(host) = nginx_matches.value_of("host") {
            // Default 8080
            let port = nginx_matches.value_of("port").unwrap_or("8080");
            breezyvps::configure::install_nginx(host);
            breezyvps::configure::add_nginx_host(host, port);
            breezyvps::configure::install_letsencrypt_cert(host);
        } else {
            println!("Missing required host parameter!");
        }
        return;
    }
}

fn main() {
    // Configure logging with simplelogger
    CombinedLogger::init(
        vec![
            TermLogger::new(LogLevelFilter::Info, Config::default()).unwrap(),
            WriteLogger::new(LogLevelFilter::Info, Config::default(), File::create("breezyvps.log").unwrap()),
        ]
    ).unwrap();

    let matches = clap_app!(myapp =>
        (version: "0.1.3")
        (author: "Michael Xu <michaeljxu11@gmail.com>")
        (about: "One stop shop for common command line goodness")
        (@arg verbose: -v ... "Enable verbose output")
        (@subcommand doctl =>
            (about: "Doctl wrapper")
            (author: "Michael Xu <michaeljxu11@gmail.com>")
            (@arg verbose: -v --verbose "Print test information verbosely")
            (@subcommand create_droplet =>
                (about: "Create a new droplet")
                (@arg name: +required "Name of the droplet, must be unique")
                (@arg config: -c "Which configuration to use [small|medium|large]")
            )
            (@subcommand destroy_droplet =>
                (about: "Destroy a droplet by name")
                (@arg name: +required "Name of the droplet to destroy completely")
            )
            (@subcommand create_sshkey =>
                (about: "Add an ssh key, which will be added upon instance creation")
                (@arg name: +required "Name of the new ssh keys")
            )
        )
        (@subcommand configure =>
            (about: "Configure droplets / nodes")
            (@subcommand nginx =>
                (about: "Install nginx and configure for host")
                (@arg host: +required "Host name of the droplet")
                (@arg port: "Port to run webapp from (default: 8080)")
            )
        )
    ).get_matches();

//    // Vary the output based on how many times the user used the "verbose" flag
//    // (i.e. 'myprog -v -v -v' or 'myprog -vvv' vs 'myprog -v'
//    match matches.occurrences_of("v") {
//        0 => println!("No verbose info"),
//        1 => println!("Some verbose info"),
//        2 => println!("Tons of verbose info"),
//        3 | _ => println!("Don't be crazy"),
//    }

    // You can handle information about subcommands by requesting their matches by name
    // (as below), requesting just the name used, or both at the same time
    if let Some(matches) = matches.subcommand_matches("doctl") {
        sc_doctl(matches);
        return;
    }
    if let Some(matches) = matches.subcommand_matches("configure") {
        sc_configure(matches);
        return;
    }

    let mut x = breezyvps::chain::CommandChain::new();
    x.cmd("echo hello")
        .cmd("test2")
        .cmd("test3")
        .execute();

    breezyvps::chain::CommandChain::new()
        .cmd_nonfatal("echo hello")
        .cmd_nonfatal("restinpeace")
        .cmd("echo yo")
        .execute();

    let processing_func = |res: &command::Result| -> command::Result {
        let mut extra_stdout = res.stdout.clone();
        extra_stdout.push_str(" + processing");

        command::Result {
            exit_code: res.exit_code,
            success: res.success,
            stdout: extra_stdout,
            stderr: res.stderr.clone()
        }
    };

    let res = breezyvps::chain::CommandChain::new()
        .cmd("echo hello")
        .result_proc(&processing_func)
        .execute();

    println!("{}", res.result.unwrap().stdout);
}
<commit_msg>slightly better handling<commit_after>#[macro_use]
extern crate clap;
extern crate breezyvps;
extern crate simplelog;

use simplelog::{Config, TermLogger, WriteLogger, CombinedLogger, LogLevelFilter};
use std::fs::File;
use breezyvps::command;

fn sc_doctl(doctl_matches: &clap::ArgMatches) {
    if let Some(create_droplet_matches) = doctl_matches.subcommand_matches("create_droplet") {
        if let Some(name) = create_droplet_matches.value_of("name") {
            breezyvps::digitalocean::create_droplet_by_name(name);
        } else {
            println!("Missing required name parameter!");
        }
        return;
    }
    if let Some(destroy_droplet_matches) = doctl_matches.subcommand_matches("destroy_droplet") {
        if let Some(name) = destroy_droplet_matches.value_of("name") {
            breezyvps::digitalocean::destroy_droplet_by_name(name);
        } else {
            println!("Missing required name parameter!");
        }
        return;
    }
    if let Some(create_ssh_key_matches) = doctl_matches.subcommand_matches("create_sshkey") {
        if let Some(name) = create_ssh_key_matches.value_of("name") {
            breezyvps::digitalocean::create_sshkey(name);
        } else {
            println!("Missing required sshkey name parameter!");
        }
        return;
    }
}

fn sc_configure(configure_matches: &clap::ArgMatches) {
    if let Some(nginx_matches) = configure_matches.subcommand_matches("nginx") {
        if let Some(host) = nginx_matches.value_of("host") {
            // Default 8080
            let port = nginx_matches.value_of("port").unwrap_or("8080");
            breezyvps::configure::install_nginx(host);
            breezyvps::configure::add_nginx_host(host, port);
            breezyvps::configure::install_letsencrypt_cert(host);
        } else {
            println!("Missing required host parameter!");
        }
        return;
    }
}

fn main() {
    // Configure logging with simplelogger
    CombinedLogger::init(
        vec![
            TermLogger::new(LogLevelFilter::Info, Config::default()).unwrap(),
            WriteLogger::new(LogLevelFilter::Info, Config::default(), File::create("breezyvps.log").unwrap()),
        ]
    ).unwrap();

    let matches = clap_app!(myapp =>
        (version: "0.1.3")
        (author: "Michael Xu <michaeljxu11@gmail.com>")
        (about: "One stop shop for common command line goodness")
        (@arg verbose: -v ... "Enable verbose output")
        (@subcommand doctl =>
            (about: "Doctl wrapper")
            (author: "Michael Xu <michaeljxu11@gmail.com>")
            (@arg verbose: -v --verbose "Print test information verbosely")
            (@subcommand create_droplet =>
                (about: "Create a new droplet")
                (@arg name: +required "Name of the droplet, must be unique")
                (@arg config: -c "Which configuration to use [small|medium|large]")
            )
            (@subcommand destroy_droplet =>
                (about: "Destroy a droplet by name")
                (@arg name: +required "Name of the droplet to destroy completely")
            )
            (@subcommand create_sshkey =>
                (about: "Add an ssh key, which will be added upon instance creation")
                (@arg name: +required "Name of the new ssh keys")
            )
        )
        (@subcommand configure =>
            (about: "Configure droplets / nodes")
            (@subcommand nginx =>
                (about: "Install nginx and configure for host")
                (@arg host: +required "Host name of the droplet")
                (@arg port: "Port to run webapp from (default: 8080)")
            )
        )
    ).get_matches();

//    // Vary the output based on how many times the user used the "verbose" flag
//    // (i.e. 'myprog -v -v -v' or 'myprog -vvv' vs 'myprog -v'
//    match matches.occurrences_of("v") {
//        0 => println!("No verbose info"),
//        1 => println!("Some verbose info"),
//        2 => println!("Tons of verbose info"),
//        3 | _ => println!("Don't be crazy"),
//    }

    // You can handle information about subcommands by requesting their matches by name
    // (as below), requesting just the name used, or both at the same time
    if let Some(matches) = matches.subcommand_matches("doctl") {
        sc_doctl(matches);
        return;
    }
    if let Some(matches) = matches.subcommand_matches("configure") {
        sc_configure(matches);
        return;
    }

    let x = breezyvps::chain::CommandChain::new();
    x.cmd("echo hello")
        .cmd("test2")
        .cmd("test3")
        .execute();

    breezyvps::chain::CommandChain::new()
        .cmd_nonfatal("echo hello")
        .cmd_nonfatal("restinpeace")
        .cmd("echo yo")
        .execute();

    let processing_func = |res: &command::Result| -> command::Result {
        let mut extra_stdout = res.stdout.clone();
        extra_stdout.push_str(" + processing");

        command::Result {
            exit_code: res.exit_code,
            success: res.success,
            stdout: extra_stdout,
            stderr: res.stderr.clone()
        }
    };

    let res = breezyvps::chain::CommandChain::new()
        .cmd("echo hello")
        .result_proc(&processing_func)
        .execute();

    println!("{}", res.result.unwrap().stdout);
}
<|endoftext|>