<commit_before>extern crate rustc_serialize;
extern crate docopt;
extern crate num;
extern crate flate2;
extern crate byteorder;
extern crate bincode;
extern crate rust_htslib as htslib;

mod tallyrun;
mod tallyread;
mod seqtable;
mod fasta;
mod filter;
mod counts;
mod bigwig;
mod scale;

use docopt::Docopt;
use std::process::exit;
use std::error::Error;
use std::path::Path;
use std::fs;
use std::fs::File;
use std::io::ErrorKind;
use seqtable::SeqTable;

/* Main usage/arguments */

const USAGE: &'static str = "
Cut-site frequencies

Usage:
  enzcut tallymer <fasta-file> <read-size> [--parts=<n>]
  enzcut seqtable <fasta-file> [options]
  enzcut dump <seqtbl-file> [<seqrange>]
  enzcut table <seqtbl-file> [<bam-file>...] [--qual=<q>] [--regions=<bedfile>] [--pdist=<min:max>] [--only-paired] [--exact-length]
  enzcut scale <seqtbl-file> <bam-file>... [options]
  enzcut <fasta-file> <bam-file>... [options]
  enzcut (-h | --help)
  enzcut --version

Options:
  -h --help             Show this screen.
  --version             Show version.
  --cut-size=<n>        Cut-site size [default: 4].
  --tallymer=<file>     Unmappable positions file produced by tallymer (seq, pos).
  --plus-offset=<p>     Cut-site offset on plus strand [default: 2]. Eg, p=2 AA[A]A.
  --minus-offset=<m>    Cut-site offset on minus strand [default: 2]. Eg, m=2 A[A]AA.
  --read-size=<r>       Read length [default: 36].
  --parts=<n>           Split suffix tree generation into n parts [default: 4].
  --qual=<q>            Minimum read quality [default: 0].
  --regions=<bedfile>   Count only cut-sites inside the regions indicated in the BED file.
  --out=<outfile>       Output seqtable filename (defaults to fasta file basename with .tbl extension).
  --bed=<bedfile>       Output scaled BED filename (defaults to BAM file basename with '_scaled.bed' extension).
  --skip-bed            Skip creating the BED file output.
  --bw=<bigwigfile>     Output scaled BigWig filename (defaults to BAM file basename with .bw extension).
  --skip-bw             Skip creating the BigWig file output.
  --stranded            Output per strand counts when writting scaled values.
  --shift-counts        Shift minus strand counts.
  --no-scale            Skip actual scalling in 'scale' command.
  --pdist=<min:max>     Distance range for included paired reads.
  --only-paired         Only accept aligned reads that have a mapped pair.
  --exact-length        Only accept BAM reads with length equal to 'read-size'.
";

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_fasta_file: String,
    arg_read_size: u16,
    arg_seqtbl_file: String,
    arg_seqrange: Option<String>,
    arg_bam_file: Option<Vec<String>>,
    flag_cut_size: u8,
    flag_tallymer: Option<String>,
    flag_plus_offset: u8,
    flag_minus_offset: u8,
    flag_version: bool,
    flag_read_size: u16,
    flag_parts: u8,
    flag_qual: u8,
    flag_regions: Option<String>,
    flag_out: Option<String>,
    flag_stranded: bool,
    flag_bed: Option<String>,
    flag_skip_bed: bool,
    flag_bw: Option<String>,
    flag_skip_bw: bool,
    flag_shift_counts: bool,
    flag_no_scale: bool,
    flag_pdist: Option<String>,
    flag_only_paired: bool,
    flag_exact_length: bool,
    cmd_tallymer: bool,
    cmd_seqtable: bool,
    cmd_dump: bool,
    cmd_table: bool,
    cmd_scale: bool,
}

fn parse_range(range: &str) -> (i32, i32) {
    let parts: Vec<&str> = range.split(':').collect();

    if parts.len() != 2 {
        println!("Invalid distance range: {}", range);
        exit(1);
    }
    let min_val = parts[0].parse::<i32>();
    let max_val = parts[1].parse::<i32>();

    if min_val.is_err() || max_val.is_err() {
        println!("Invalid distance range: {}", range);
        exit(1);
    }
    (min_val.unwrap(),max_val.unwrap())
}

fn stem_filename(stem_src: &str, suffix: &str, out_arg: Option<String>) -> String {
    match out_arg {
        Some(out) => out,
        None => {
            let mut basename = Path::new(stem_src).file_stem().unwrap().to_os_string();
            basename.push(suffix);
            basename.into_string().unwrap()
        }
    }
}

fn file_exists(filename: &str) -> bool {
	match fs::metadata(filename) {
		Ok(meta) => meta.is_file(),
		Err(err) => match err.kind() {
			ErrorKind::NotFound => false,
			_ => { println!("Failed to open file {}: {:}", filename, err); exit(-1) },
		},
	}
}

fn main() {
    // Parse command line arguments
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    if args.flag_version {
        println!("Cut-site frequencies, v{}.{}.{}",
            env!( "CARGO_PKG_VERSION_MAJOR" ),
            env!( "CARGO_PKG_VERSION_MINOR" ),
            env!( "CARGO_PKG_VERSION_PATCH" ) );
        return;
    }

    let dist_range = match args.flag_pdist {
        Some(range) => Some(parse_range(&range)),
        None => None
    };

    // Check for data output commands

    if args.cmd_dump {
        if let Some(seqrange) = args.arg_seqrange {
            seqtable::dump_seqtable_range(&args.arg_seqtbl_file, &seqrange);
        } else {
            seqtable::dump_seqtable(&args.arg_seqtbl_file);
        }

        return;
    }

    if args.cmd_table {
        let has_bam = args.arg_bam_file.is_some();
        let counts = counts::tabulate(&args.arg_seqtbl_file, args.arg_bam_file.as_ref(), args.flag_qual, args.flag_regions, dist_range, args.flag_only_paired, args.flag_exact_length);
        counts::print_counts(&counts, has_bam);
        return;
    }

    // Check for main sequence commands
    let mut run_tallymer = true;
    let mut run_seqtable = true;
    let mut run_scale = true;

    if args.cmd_tallymer {
        // just phase 1
        run_seqtable = false;
        run_scale = false;
    } else if args.cmd_seqtable {
        // phase 1 & 2
        run_scale = false;
    } else if args.cmd_scale {
        // just phase 3
        run_tallymer = false;
        run_seqtable = false;
    } else {
        // all three phases

        // catch cmd names being interpreted as fasta_file names
        if args.arg_fasta_file.eq("dump") || args.arg_fasta_file.eq("table") || args.arg_fasta_file.eq("tallymer") || args.arg_fasta_file.eq("seqtable") || args.arg_fasta_file.eq("scale") {
            println!("Invalid arguments to {} command.", args.arg_fasta_file);
            println!("{}", USAGE);
            exit(1);
        }
    }

    if args.flag_skip_bed && args.flag_skip_bw {
        println!("--skip-bed and --skip-bw cannot be used together");
        exit(1);
    }

    // Process main sequence phases

    // phase 1 - tallymer
    let tally_path = if run_tallymer {
        if !file_exists(&args.arg_fasta_file) {
            println!("Error: FASTA file {} does not exist!", args.arg_fasta_file);
            exit(1);
        }

        let path = tallyrun::tallymer_createfile(&args.arg_fasta_file, args.flag_read_size, args.flag_parts);
        println!("# tallymer produced/found {:}", path.to_str().unwrap());
        Some(path)
    } else { None };

    // phase 2 - seqtable
    let seqtable_file = if run_seqtable {
        let suffix = format!("_{}.{}.{}.{}.tbl", args.flag_read_size, args.flag_cut_size, args.flag_plus_offset, args.flag_minus_offset);
        let outfile = stem_filename(&args.arg_fasta_file, &suffix, args.flag_out);

        let seq_params = seqtable::SeqTableParams {
            cut_length: args.flag_cut_size,
            plus_offset: args.flag_plus_offset,
            minus_offset: args.flag_minus_offset,
            read_length: args.flag_read_size };

        if file_exists(&outfile) {
            let file = File::open(&outfile).ok().expect("read file");
            let table = match SeqTable::open(file) {
                Ok(value) => value,
                Err(e) => {
                    println!("Error: seqtable: {}", e.to_string());
                    exit(1);
                },
            };
            if table.equivalent(&outfile, &seq_params) {
                println!("# seqtable reusing existing {}", &outfile);
                outfile
            } else {
                println!("Error: seqtable: output file {} already exists but does not match requested parameters!", outfile);
                exit(1);
            }
        } else {
            fasta::process_fasta(&args.arg_fasta_file, &tally_path.unwrap(), seq_params, &outfile);
            println!("# seqtable produced {}", &outfile);
            outfile
        }
    } else {
        args.arg_seqtbl_file
    };

    // phase 3 - tabulate & scale
    if run_scale {
        for filename in args.arg_bam_file.as_ref().unwrap() {
            if !file_exists(filename) {
                println!("Error: BAM file {} does not exist!", filename);
                exit(1);
            }
        }

        let bamfile = args.arg_bam_file.as_ref().unwrap()[0].clone(); // use the first name for reference
        let counts = counts::tabulate(&seqtable_file, args.arg_bam_file.as_ref(), args.flag_qual, args.flag_regions, dist_range, args.flag_only_paired, args.flag_exact_length);
        let pileup = scale::scale(&seqtable_file, counts, args.arg_bam_file.as_ref().unwrap(), args.flag_qual, args.flag_shift_counts, args.flag_no_scale, &dist_range, args.flag_only_paired, args.flag_exact_length);

        if !args.flag_skip_bed {
            let outfile_bed = if args.flag_no_scale {
                    stem_filename(&bamfile, "_not_scaled.bed", args.flag_bed)
                } else {
                    stem_filename(&bamfile, "_scaled.bed", args.flag_bed)
                };

            if file_exists(&outfile_bed) {
                println!("Error: output BED file {} already exists!", outfile_bed);
                exit(1);
            }

            match pileup.write_bed(&outfile_bed, args.flag_stranded) {
                Ok(_) => println!("# scale produced {}", &outfile_bed),
                Err(err) => println!("Error producing BED file: {}", err.description()),
            }
        } else {
            println!("# scale skipping BED output");
        }

        if !args.flag_skip_bw {
            let outfile_bw = stem_filename(&bamfile, ".bw", args.flag_bw);

            match pileup.write_bw(&outfile_bw, args.flag_stranded) {
                Ok((f1, f2)) => {
                    if f2.is_some() {
                        println!("# scale produced {}", f1);
                        println!("# scale produced {}", f2.unwrap());
                    } else {
                        println!("# scale produced {}", f1);
                    }
                },
                Err(err) => println!("Error producing BigWig file: {}", err.description()),
            }
        } else {
            println!("# scale skipping BigWig output");
        }
    }
}
<commit_msg>Fix docopt default issue<commit_after>extern crate rustc_serialize;
extern crate docopt;
extern crate num;
extern crate flate2;
extern crate byteorder;
extern crate bincode;
extern crate rust_htslib as htslib;

mod tallyrun;
mod tallyread;
mod seqtable;
mod fasta;
mod filter;
mod counts;
mod bigwig;
mod scale;

use docopt::Docopt;
use std::process::exit;
use std::error::Error;
use std::path::Path;
use std::fs;
use std::fs::File;
use std::io::ErrorKind;
use seqtable::SeqTable;

/* Main usage/arguments */

const USAGE: &'static str = "
Cut-site frequencies

Usage:
  enzcut tallymer <fasta-file> <read-size> [--parts=<n>]
  enzcut seqtable <fasta-file> [options]
  enzcut dump <seqtbl-file> [<seqrange>]
  enzcut table <seqtbl-file> [<bam-file>...] [--qual=<q>] [--regions=<bedfile>] [--pdist=<min:max>] [--only-paired] [--exact-length]
  enzcut scale <seqtbl-file> <bam-file>... [options]
  enzcut <fasta-file> <bam-file>... [options]
  enzcut (-h | --help)
  enzcut --version

Options:
  -h --help             Show this screen.
  --version             Show version.
  --cut-size=<n>        Cut-site size [default: 4].
  --tallymer=<file>     Unmappable positions file produced by tallymer (seq, pos).
  --plus-offset=<p>     Cut-site offset on plus strand, eg. p=2 AA[A]A [default: 2].
  --minus-offset=<m>    Cut-site offset on minus strand, eg. Eg, m=2 A[A]AA [default: 2].
  --read-size=<r>       Read length [default: 36].
  --parts=<n>           Split suffix tree generation into n parts [default: 4].
  --qual=<q>            Minimum read quality [default: 0].
  --regions=<bedfile>   Count only cut-sites inside the regions indicated in the BED file.
  --out=<outfile>       Output seqtable filename (defaults to fasta file basename with .tbl extension).
  --bed=<bedfile>       Output scaled BED filename (defaults to BAM file basename with '_scaled.bed' extension).
  --skip-bed            Skip creating the BED file output.
  --bw=<bigwigfile>     Output scaled BigWig filename (defaults to BAM file basename with .bw extension).
  --skip-bw             Skip creating the BigWig file output.
  --stranded            Output per strand counts when writting scaled values.
  --shift-counts        Shift minus strand counts.
  --no-scale            Skip actual scalling in 'scale' command.
  --pdist=<min:max>     Distance range for included paired reads.
  --only-paired         Only accept aligned reads that have a mapped pair.
  --exact-length        Only accept BAM reads with length equal to 'read-size'.
";

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_fasta_file: String,
    arg_read_size: u16,
    arg_seqtbl_file: String,
    arg_seqrange: Option<String>,
    arg_bam_file: Option<Vec<String>>,
    flag_cut_size: u8,
    flag_tallymer: Option<String>,
    flag_plus_offset: u8,
    flag_minus_offset: u8,
    flag_version: bool,
    flag_read_size: u16,
    flag_parts: u8,
    flag_qual: u8,
    flag_regions: Option<String>,
    flag_out: Option<String>,
    flag_stranded: bool,
    flag_bed: Option<String>,
    flag_skip_bed: bool,
    flag_bw: Option<String>,
    flag_skip_bw: bool,
    flag_shift_counts: bool,
    flag_no_scale: bool,
    flag_pdist: Option<String>,
    flag_only_paired: bool,
    flag_exact_length: bool,
    cmd_tallymer: bool,
    cmd_seqtable: bool,
    cmd_dump: bool,
    cmd_table: bool,
    cmd_scale: bool,
}

fn parse_range(range: &str) -> (i32, i32) {
    let parts: Vec<&str> = range.split(':').collect();

    if parts.len() != 2 {
        println!("Invalid distance range: {}", range);
        exit(1);
    }
    let min_val = parts[0].parse::<i32>();
    let max_val = parts[1].parse::<i32>();

    if min_val.is_err() || max_val.is_err() {
        println!("Invalid distance range: {}", range);
        exit(1);
    }
    (min_val.unwrap(),max_val.unwrap())
}

fn stem_filename(stem_src: &str, suffix: &str, out_arg: Option<String>) -> String {
    match out_arg {
        Some(out) => out,
        None => {
            let mut basename = Path::new(stem_src).file_stem().unwrap().to_os_string();
            basename.push(suffix);
            basename.into_string().unwrap()
        }
    }
}

fn file_exists(filename: &str) -> bool {
	match fs::metadata(filename) {
		Ok(meta) => meta.is_file(),
		Err(err) => match err.kind() {
			ErrorKind::NotFound => false,
			_ => { println!("Failed to open file {}: {:}", filename, err); exit(-1) },
		},
	}
}

fn main() {
    // Parse command line arguments
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    if args.flag_version {
        println!("Cut-site frequencies, v{}.{}.{}",
            env!( "CARGO_PKG_VERSION_MAJOR" ),
            env!( "CARGO_PKG_VERSION_MINOR" ),
            env!( "CARGO_PKG_VERSION_PATCH" ) );
        return;
    }

    let dist_range = match args.flag_pdist {
        Some(range) => Some(parse_range(&range)),
        None => None
    };

    // Check for data output commands

    if args.cmd_dump {
        if let Some(seqrange) = args.arg_seqrange {
            seqtable::dump_seqtable_range(&args.arg_seqtbl_file, &seqrange);
        } else {
            seqtable::dump_seqtable(&args.arg_seqtbl_file);
        }

        return;
    }

    if args.cmd_table {
        let has_bam = args.arg_bam_file.is_some();
        let counts = counts::tabulate(&args.arg_seqtbl_file, args.arg_bam_file.as_ref(), args.flag_qual, args.flag_regions, dist_range, args.flag_only_paired, args.flag_exact_length);
        counts::print_counts(&counts, has_bam);
        return;
    }

    // Check for main sequence commands
    let mut run_tallymer = true;
    let mut run_seqtable = true;
    let mut run_scale = true;

    if args.cmd_tallymer {
        // just phase 1
        run_seqtable = false;
        run_scale = false;
    } else if args.cmd_seqtable {
        // phase 1 & 2
        run_scale = false;
    } else if args.cmd_scale {
        // just phase 3
        run_tallymer = false;
        run_seqtable = false;
    } else {
        // all three phases

        // catch cmd names being interpreted as fasta_file names
        if args.arg_fasta_file.eq("dump") || args.arg_fasta_file.eq("table") || args.arg_fasta_file.eq("tallymer") || args.arg_fasta_file.eq("seqtable") || args.arg_fasta_file.eq("scale") {
            println!("Invalid arguments to {} command.", args.arg_fasta_file);
            println!("{}", USAGE);
            exit(1);
        }
    }

    if args.flag_skip_bed && args.flag_skip_bw {
        println!("--skip-bed and --skip-bw cannot be used together");
        exit(1);
    }

    // Process main sequence phases

    // phase 1 - tallymer
    let tally_path = if run_tallymer {
        if !file_exists(&args.arg_fasta_file) {
            println!("Error: FASTA file {} does not exist!", args.arg_fasta_file);
            exit(1);
        }

        let path = tallyrun::tallymer_createfile(&args.arg_fasta_file, args.flag_read_size, args.flag_parts);
        println!("# tallymer produced/found {:}", path.to_str().unwrap());
        Some(path)
    } else { None };

    // phase 2 - seqtable
    let seqtable_file = if run_seqtable {
        let suffix = format!("_{}.{}.{}.{}.tbl", args.flag_read_size, args.flag_cut_size, args.flag_plus_offset, args.flag_minus_offset);
        let outfile = stem_filename(&args.arg_fasta_file, &suffix, args.flag_out);

        let seq_params = seqtable::SeqTableParams {
            cut_length: args.flag_cut_size,
            plus_offset: args.flag_plus_offset,
            minus_offset: args.flag_minus_offset,
            read_length: args.flag_read_size };

        if file_exists(&outfile) {
            let file = File::open(&outfile).ok().expect("read file");
            let table = match SeqTable::open(file) {
                Ok(value) => value,
                Err(e) => {
                    println!("Error: seqtable: {}", e.to_string());
                    exit(1);
                },
            };
            if table.equivalent(&outfile, &seq_params) {
                println!("# seqtable reusing existing {}", &outfile);
                outfile
            } else {
                println!("Error: seqtable: output file {} already exists but does not match requested parameters!", outfile);
                exit(1);
            }
        } else {
            fasta::process_fasta(&args.arg_fasta_file, &tally_path.unwrap(), seq_params, &outfile);
            println!("# seqtable produced {}", &outfile);
            outfile
        }
    } else {
        args.arg_seqtbl_file
    };

    // phase 3 - tabulate & scale
    if run_scale {
        for filename in args.arg_bam_file.as_ref().unwrap() {
            if !file_exists(filename) {
                println!("Error: BAM file {} does not exist!", filename);
                exit(1);
            }
        }

        let bamfile = args.arg_bam_file.as_ref().unwrap()[0].clone(); // use the first name for reference
        let counts = counts::tabulate(&seqtable_file, args.arg_bam_file.as_ref(), args.flag_qual, args.flag_regions, dist_range, args.flag_only_paired, args.flag_exact_length);
        let pileup = scale::scale(&seqtable_file, counts, args.arg_bam_file.as_ref().unwrap(), args.flag_qual, args.flag_shift_counts, args.flag_no_scale, &dist_range, args.flag_only_paired, args.flag_exact_length);

        if !args.flag_skip_bed {
            let outfile_bed = if args.flag_no_scale {
                    stem_filename(&bamfile, "_not_scaled.bed", args.flag_bed)
                } else {
                    stem_filename(&bamfile, "_scaled.bed", args.flag_bed)
                };

            if file_exists(&outfile_bed) {
                println!("Error: output BED file {} already exists!", outfile_bed);
                exit(1);
            }

            match pileup.write_bed(&outfile_bed, args.flag_stranded) {
                Ok(_) => println!("# scale produced {}", &outfile_bed),
                Err(err) => println!("Error producing BED file: {}", err.description()),
            }
        } else {
            println!("# scale skipping BED output");
        }

        if !args.flag_skip_bw {
            let outfile_bw = stem_filename(&bamfile, ".bw", args.flag_bw);

            match pileup.write_bw(&outfile_bw, args.flag_stranded) {
                Ok((f1, f2)) => {
                    if f2.is_some() {
                        println!("# scale produced {}", f1);
                        println!("# scale produced {}", f2.unwrap());
                    } else {
                        println!("# scale produced {}", f1);
                    }
                },
                Err(err) => println!("Error producing BigWig file: {}", err.description()),
            }
        } else {
            println!("# scale skipping BigWig output");
        }
    }
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#![feature(iter_arith, zero_one, custom_derive,
           custom_attribute, plugin, const_fn)]
#![plugin(serde_macros)]
#![plugin(clippy)]

#![allow(match_same_arms)] // we seem to have instances where same arms are good
#![allow(dead_code)] // only temporary, until more stuff is filled in

extern crate devicemapper;
#[macro_use] extern crate clap;
#[macro_use] extern crate nix;
extern crate crc;
extern crate byteorder;
extern crate uuid;
extern crate time;
extern crate serde;
extern crate serde_json;
extern crate bytesize;
extern crate dbus;
extern crate term;

#[macro_use] extern crate custom_derive;
#[macro_use] extern crate newtype_derive;

pub static mut debug: bool = false;

macro_rules! dbgp {
    ($($arg:tt)*) => (
        unsafe {
            if ::debug {
                println!($($arg)*)
            }
        })
}

mod types;
mod consts;
mod froyo;
mod blockdev;
mod raid;
mod thin;
mod util;
mod dbus_api;

use std::io::Write;
use std::error::Error;
use std::process::exit;
use std::path::{Path, PathBuf};
use std::rc::Rc;
use std::cell::RefCell;

use clap::{App, Arg, SubCommand, ArgMatches};
use bytesize::ByteSize;
use dbus::{Connection, BusType};

use types::FroyoResult;
use consts::SECTOR_SIZE;
use froyo::{Froyo, FroyoStatus};

// We are given BlockDevs to start.
// We allocate LinearDevs from each for the meta and data devices.
// We use all these to make RaidDevs.
// We create two RaidLinearDevs from these for meta and data devices.
// We use these to make a ThinPoolDev.
// From that, we allocate a ThinDev.

fn list(_args: &ArgMatches) -> FroyoResult<()> {
    let froyos = try!(Froyo::find_all());
    for f in &froyos {
        println!("{}", f.name);
    }

    Ok(())
}

fn status(args: &ArgMatches) -> FroyoResult<()> {
    let name = args.value_of("froyodevname").unwrap();
    match try!(Froyo::find(&name)) {
        Some(f) => {
            let status = match try!(f.status()) {
                FroyoStatus::Good(_) => "good",
                FroyoStatus::RaidFailed => "RAID failed",
                FroyoStatus::ThinPoolFailed => "Thin pool failed",
                FroyoStatus::ThinFailed => "Thin device failed",
            };

            let space = *try!(f.avail_redundant_space()) * SECTOR_SIZE;
            let total = *f.total_redundant_space() * SECTOR_SIZE;

            let percent = ((total-space) * 100) / total;

            println!("Status: {}, {}% used ({} of {} free)",
                     status, percent,
                     ByteSize::b(space as usize).to_string(true),
                     ByteSize::b(total as usize).to_string(true));
        },
        None => println!("Froyodev \"{}\" not found", name),
    }
    Ok(())
}

fn add(_args: &ArgMatches) -> FroyoResult<()> {
    println!("hello from add()");
    Ok(())
}

fn remove(_args: &ArgMatches) -> FroyoResult<()> {
    let mut out = term::stdout().expect("could not get stdout");

    try!(out.fg(term::color::GREEN));
    try!(writeln!(out, "hello color"));
    try!(out.reset());

    Ok(())
}

fn create(args: &ArgMatches) -> FroyoResult<()> {
    let name = args.value_of("froyodevname").unwrap();
    let dev_paths: Vec<_> = args.values_of("devices").unwrap().into_iter()
        .map(|dev| {
            if !Path::new(dev).is_absolute() {
                PathBuf::from(format!("/dev/{}", dev))
            } else {
                PathBuf::from(dev)
            }})
        .collect();
    let force = args.is_present("force");

    let froyo = try!(Froyo::new(name, &dev_paths, force));

    try!(froyo.save_state());

    dbgp!("Froyodev {} created", froyo.name);

    Ok(())
}

fn dump_meta(args: &ArgMatches) -> FroyoResult<()> {
    let name = args.value_of("froyodevname").unwrap();
    match try!(Froyo::find(&name)) {
        Some(f) =>
            println!("{}", try!(f.to_metadata_pretty())),
        None => println!("Froyodev \"{}\" not found", name),
    }

    Ok(())
}

fn dbus_server(_args: &ArgMatches) -> FroyoResult<()> {
    let c = Connection::get_private(BusType::Session).unwrap();
    let froyos = try!(Froyo::find_all());
    let froyos = froyos.into_iter()
        .map(|f| Rc::new(RefCell::new(f)))
        .collect::<Vec<_>>();
    let mut froyos = Rc::new(RefCell::new(froyos));
    let tree = try!(dbus_api::get_tree(&c, &mut froyos));

    // TODO: event loop needs to handle dbus and also dm events (or polling)
    // so we can extend/reshape/delay/whatever in a timely fashion
    for _ in tree.run(&c, c.iter(1000)) {

    }

    Ok(())
}

fn teardown(args: &ArgMatches) -> FroyoResult<()> {
    let name = args.value_of("froyodevname").unwrap();
    match try!(Froyo::find(&name)) {
        Some(mut f) => try!(f.teardown()),
        None => println!("Froyodev \"{}\" not found", name),
    }

    Ok(())
}

fn main() {

    let matches = App::new("froyo")
        .version(&crate_version!())
        .author("Andy Grover <andy@groveronline.com>")
        .about("Drobo + Free + YOLO")
        .arg(Arg::with_name("debug")
             .short("d")
             .long("debug")
             .help("Print additional output for debugging")
             )
        .subcommand(SubCommand::with_name("list")
                    .about("List all froyodevs")
                    .arg(Arg::with_name("long")
                         .short("l")
                         .long("long")
                         .help("Use a long listing format")
                         )
                    )
        .subcommand(SubCommand::with_name("status")
                    .about("Get the status of a single froyodev")
                    .arg(Arg::with_name("froyodevname")
                         .required(true)
                         .help("Froyodev to get info on")
                         .index(1)
                         )
                    )
        .subcommand(SubCommand::with_name("add")
                    .about("Add one or more additional block devices to a froyodev")
                    .arg(Arg::with_name("froyodevname")
                         .help("Froyodev to add the device to")
                         .required(true)
                         .index(1)
                         )
                    .arg(Arg::with_name("devices")
                         .help("device(s) to add")
                         .multiple(true)
                         .required(true)
                         .index(2)
                         )
                    )
        .subcommand(SubCommand::with_name("remove")
                    .about("Remove a block device from a froyodev")
                    .arg(Arg::with_name("froyodevname")
                         .help("Froyodev to remove the device from")
                         .required(true)
                         .index(1)
                         )
                    .arg(Arg::with_name("devices")
                         .help("Block device(s) to remove")
                         .multiple(true)
                         .required(true)
                         .index(2)
                         )
                    )
        .subcommand(SubCommand::with_name("create")
                    .about("Create a new froyodev")
                    .arg(Arg::with_name("force")
                         .short("f")
                         .long("force")
                         .help("Force")
                         )
                    .arg(Arg::with_name("froyodevname")
                         .help("Name of the new froyodev")
                         .required(true)
                         .index(1)
                         )
                    .arg(Arg::with_name("devices")
                         .help("Initial block device(s) to use")
                         .multiple(true)
                         .required(true)
                         .index(2)
                         )
                    )
        .subcommand(SubCommand::with_name("dev")
                    .about("Developer/debug commands")
                    .subcommand(SubCommand::with_name("dump_meta")
                                .about("Output the JSON metadata for a froyodev")
                                .arg(Arg::with_name("froyodevname")
                                     .help("Name of the froyodev")
                                     .required(true)
                                     .index(1)
                                     )
                                )
                    .subcommand(SubCommand::with_name("teardown")
                                .about("Remove the DM mappings for a froyodev")
                                .arg(Arg::with_name("froyodevname")
                                     .help("Name of the froyodev")
                                     .required(true)
                                     .index(1)
                                     )
                                )
                    .subcommand(SubCommand::with_name("dbus_server")
                                .about("Serve the Froyo DBus API")
                                )
                    )
        .get_matches();

    if matches.is_present("debug") {
        // must use unsafe to change a mut static, sigh
        unsafe { debug = true };
    }

    let r = match matches.subcommand() {
        ("list", Some(matches)) => list(matches),
        ("status", Some(matches)) => status(matches),
        ("add", Some(matches)) => add(matches),
        ("remove", Some(matches)) => remove(matches),
        ("create", Some(matches)) => create(matches),
        ("dev", Some(matches)) => match matches.subcommand() {
            ("dump_meta", Some(matches)) => dump_meta(matches),
            ("dbus_server", Some(matches)) => dbus_server(matches),
            ("teardown", Some(matches)) => teardown(matches),
            ("", None) => {
                println!("No command given, try \"help\"");
                Ok(())
            }
            _ => unreachable!(),
        },
        ("", None) => {
            println!("No command given, try \"help\"");
            Ok(())
        }
        _ => unreachable!(),
    };

    if let Err(r) = r {
        if let Err(e) = writeln!(&mut ::std::io::stderr(), "{}", r.description()) {
            panic!("Unable to write to stderr: {}", e)
        }
        exit(1);
    }
}
<commit_msg>main: Rough-in rename and destroy subcommands<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#![feature(iter_arith, zero_one, custom_derive,
           custom_attribute, plugin, const_fn)]
#![plugin(serde_macros)]
#![plugin(clippy)]

#![allow(match_same_arms)] // we seem to have instances where same arms are good
#![allow(dead_code)] // only temporary, until more stuff is filled in

extern crate devicemapper;
#[macro_use] extern crate clap;
#[macro_use] extern crate nix;
extern crate crc;
extern crate byteorder;
extern crate uuid;
extern crate time;
extern crate serde;
extern crate serde_json;
extern crate bytesize;
extern crate dbus;
extern crate term;

#[macro_use] extern crate custom_derive;
#[macro_use] extern crate newtype_derive;

pub static mut debug: bool = false;

macro_rules! dbgp {
    ($($arg:tt)*) => (
        unsafe {
            if ::debug {
                println!($($arg)*)
            }
        })
}

mod types;
mod consts;
mod froyo;
mod blockdev;
mod raid;
mod thin;
mod util;
mod dbus_api;

use std::io::Write;
use std::error::Error;
use std::process::exit;
use std::path::{Path, PathBuf};
use std::rc::Rc;
use std::cell::RefCell;

use clap::{App, Arg, SubCommand, ArgMatches};
use bytesize::ByteSize;
use dbus::{Connection, BusType};

use types::FroyoResult;
use consts::SECTOR_SIZE;
use froyo::{Froyo, FroyoStatus};

// We are given BlockDevs to start.
// We allocate LinearDevs from each for the meta and data devices.
// We use all these to make RaidDevs.
// We create two RaidLinearDevs from these for meta and data devices.
// We use these to make a ThinPoolDev.
// From that, we allocate a ThinDev.

fn list(_args: &ArgMatches) -> FroyoResult<()> {
    let froyos = try!(Froyo::find_all());
    for f in &froyos {
        println!("{}", f.name);
    }

    Ok(())
}

fn status(args: &ArgMatches) -> FroyoResult<()> {
    let name = args.value_of("froyodevname").unwrap();
    match try!(Froyo::find(&name)) {
        Some(f) => {
            let status = match try!(f.status()) {
                FroyoStatus::Good(_) => "good",
                FroyoStatus::RaidFailed => "RAID failed",
                FroyoStatus::ThinPoolFailed => "Thin pool failed",
                FroyoStatus::ThinFailed => "Thin device failed",
            };

            let space = *try!(f.avail_redundant_space()) * SECTOR_SIZE;
            let total = *f.total_redundant_space() * SECTOR_SIZE;

            let percent = ((total-space) * 100) / total;

            println!("Status: {}, {}% used ({} of {} free)",
                     status, percent,
                     ByteSize::b(space as usize).to_string(true),
                     ByteSize::b(total as usize).to_string(true));
        },
        None => println!("Froyodev \"{}\" not found", name),
    }
    Ok(())
}

fn add(_args: &ArgMatches) -> FroyoResult<()> {
    println!("hello from add()");
    Ok(())
}

fn remove(_args: &ArgMatches) -> FroyoResult<()> {
    let mut out = term::stdout().expect("could not get stdout");

    try!(out.fg(term::color::GREEN));
    try!(writeln!(out, "hello color"));
    try!(out.reset());

    Ok(())
}

fn create(args: &ArgMatches) -> FroyoResult<()> {
    let name = args.value_of("froyodevname").unwrap();
    let dev_paths: Vec<_> = args.values_of("devices").unwrap().into_iter()
        .map(|dev| {
            if !Path::new(dev).is_absolute() {
                PathBuf::from(format!("/dev/{}", dev))
            } else {
                PathBuf::from(dev)
            }})
        .collect();
    let force = args.is_present("force");

    let froyo = try!(Froyo::new(name, &dev_paths, force));

    try!(froyo.save_state());

    dbgp!("Froyodev {} created", froyo.name);

    Ok(())
}

fn rename(args: &ArgMatches) -> FroyoResult<()> {
    let old_name = args.value_of("froyodev_old_name").unwrap();
    let new_name = args.value_of("froyodev_new_name").unwrap();
    println!("Rename {} to {}", old_name, new_name);

    Ok(())
}

fn destroy(args: &ArgMatches) -> FroyoResult<()> {
    let name = args.value_of("froyodev").unwrap();
    println!("Destroy froyodev {}", name);

    Ok(())
}

fn dump_meta(args: &ArgMatches) -> FroyoResult<()> {
    let name = args.value_of("froyodevname").unwrap();
    match try!(Froyo::find(&name)) {
        Some(f) =>
            println!("{}", try!(f.to_metadata_pretty())),
        None => println!("Froyodev \"{}\" not found", name),
    }

    Ok(())
}

fn dbus_server(_args: &ArgMatches) -> FroyoResult<()> {
    let c = Connection::get_private(BusType::Session).unwrap();
    let froyos = try!(Froyo::find_all());
    let froyos = froyos.into_iter()
        .map(|f| Rc::new(RefCell::new(f)))
        .collect::<Vec<_>>();
    let mut froyos = Rc::new(RefCell::new(froyos));
    let tree = try!(dbus_api::get_tree(&c, &mut froyos));

    // TODO: event loop needs to handle dbus and also dm events (or polling)
    // so we can extend/reshape/delay/whatever in a timely fashion
    for _ in tree.run(&c, c.iter(1000)) {

    }

    Ok(())
}

fn teardown(args: &ArgMatches) -> FroyoResult<()> {
    let name = args.value_of("froyodevname").unwrap();
    match try!(Froyo::find(&name)) {
        Some(mut f) => try!(f.teardown()),
        None => println!("Froyodev \"{}\" not found", name),
    }

    Ok(())
}

fn main() {

    let matches = App::new("froyo")
        .version(&crate_version!())
        .author("Andy Grover <andy@groveronline.com>")
        .about("Drobo + Free + YOLO")
        .arg(Arg::with_name("debug")
             .short("d")
             .long("debug")
             .help("Print additional output for debugging")
             )
        .subcommand(SubCommand::with_name("list")
                    .about("List all froyodevs")
                    .arg(Arg::with_name("long")
                         .short("l")
                         .long("long")
                         .help("Use a long listing format")
                         )
                    )
        .subcommand(SubCommand::with_name("status")
                    .about("Get the status of a single froyodev")
                    .arg(Arg::with_name("froyodevname")
                         .required(true)
                         .help("Froyodev to get info on")
                         .index(1)
                         )
                    )
        .subcommand(SubCommand::with_name("add")
                    .about("Add one or more additional block devices to a froyodev")
                    .arg(Arg::with_name("froyodevname")
                         .help("Froyodev to add the device to")
                         .required(true)
                         .index(1)
                         )
                    .arg(Arg::with_name("devices")
                         .help("device(s) to add")
                         .multiple(true)
                         .required(true)
                         .index(2)
                         )
                    )
        .subcommand(SubCommand::with_name("remove")
                    .about("Remove a block device from a froyodev")
                    .arg(Arg::with_name("froyodevname")
                         .help("Froyodev to remove the device from")
                         .required(true)
                         .index(1)
                         )
                    .arg(Arg::with_name("devices")
                         .help("Block device(s) to remove")
                         .multiple(true)
                         .required(true)
                         .index(2)
                         )
                    )
        .subcommand(SubCommand::with_name("create")
                    .about("Create a new froyodev")
                    .arg(Arg::with_name("force")
                         .short("f")
                         .long("force")
                         .help("Force")
                         )
                    .arg(Arg::with_name("froyodevname")
                         .help("Name of the new froyodev")
                         .required(true)
                         .index(1)
                         )
                    .arg(Arg::with_name("devices")
                         .help("Initial block device(s) to use")
                         .multiple(true)
                         .required(true)
                         .index(2)
                         )
                    )
        .subcommand(SubCommand::with_name("rename")
                    .about("Rename a froyodev")
                    .arg(Arg::with_name("froyodev_old_name")
                         .help("Old name of froyodev")
                         .required(true)
                         .index(1)
                    )
                    .arg(Arg::with_name("froyodev_new_name")
                         .help("New name of froyodev")
                         .required(true)
                         .index(2)
                    )
        )
        .subcommand(SubCommand::with_name("destroy")
                    .about("Destroy a froyodev")
                    .arg(Arg::with_name("froyodev")
                         .help("Froyodev to destroy")
                         .required(true)
                         .index(1)
                    )
        )
        .subcommand(SubCommand::with_name("dev")
                    .about("Developer/debug commands")
                    .subcommand(SubCommand::with_name("dump_meta")
                                .about("Output the JSON metadata for a froyodev")
                                .arg(Arg::with_name("froyodevname")
                                     .help("Name of the froyodev")
                                     .required(true)
                                     .index(1)
                                     )
                                )
                    .subcommand(SubCommand::with_name("teardown")
                                .about("Remove the DM mappings for a froyodev")
                                .arg(Arg::with_name("froyodevname")
                                     .help("Name of the froyodev")
                                     .required(true)
                                     .index(1)
                                     )
                                )
                    .subcommand(SubCommand::with_name("dbus_server")
                                .about("Serve the Froyo DBus API")
                                )
                    )
        .get_matches();

    if matches.is_present("debug") {
        // must use unsafe to change a mut static, sigh
        unsafe { debug = true };
    }

    let r = match matches.subcommand() {
        ("list", Some(matches)) => list(matches),
        ("status", Some(matches)) => status(matches),
        ("add", Some(matches)) => add(matches),
        ("remove", Some(matches)) => remove(matches),
        ("create", Some(matches)) => create(matches),
        ("rename", Some(matches)) => rename(matches),
        ("destroy", Some(matches)) => destroy(matches),
        ("dev", Some(matches)) => match matches.subcommand() {
            ("dump_meta", Some(matches)) => dump_meta(matches),
            ("dbus_server", Some(matches)) => dbus_server(matches),
            ("teardown", Some(matches)) => teardown(matches),
            ("", None) => {
                println!("No command given, try \"help\"");
                Ok(())
            }
            _ => unreachable!(),
        },
        ("", None) => {
            println!("No command given, try \"help\"");
            Ok(())
        }
        _ => unreachable!(),
    };

    if let Err(r) = r {
        if let Err(e) = writeln!(&mut ::std::io::stderr(), "{}", r.description()) {
            panic!("Unable to write to stderr: {}", e)
        }
        exit(1);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;

use petname::Petnames;

use std::collections::HashSet;
use std::fmt;
use std::fs;
use std::io::{self, Write};
use std::path;
use std::process;
use std::str::FromStr;

use clap::Arg;
use rand::seq::IteratorRandom;

fn main() {
    let matches = app().get_matches();
    match run(matches) {
        Err(Error::Disconnected) => {
            process::exit(0);
        }
        Err(e) => {
            eprintln!("Error: {}", e);
            process::exit(1);
        }
        Ok(()) => {
            process::exit(0);
        }
    }
}

enum Error {
    Io(io::Error),
    FileIo(path::PathBuf, io::Error),
    Cardinality(String),
    Alliteration(String),
    Disconnected,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Error::Io(ref e) => write!(f, "{}", e),
            Error::FileIo(ref path, ref e) => write!(f, "{}: {}", e, path.display()),
            Error::Cardinality(ref message) => write!(f, "cardinality is zero: {}", message),
            Error::Alliteration(ref message) => write!(f, "cannot alliterate: {}", message),
            Error::Disconnected => write!(f, "caller disconnected / stopped reading"),
        }
    }
}

impl From<io::Error> for Error {
    fn from(error: io::Error) -> Self {
        Error::Io(error)
    }
}

fn app<'a, 'b>() -> clap::App<'a, 'b> {
    clap::App::new("rust-petname")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Generate human readable random names.")
        .after_help(concat!(
            "Based on Dustin Kirkland's petname project ",
            "<https://github.com/dustinkirkland/petname>."
        ))
        .arg(
            Arg::with_name("words")
                .short("w")
                .long("words")
                .value_name("WORDS")
                .default_value("2")
                .help("Number of words in name")
                .takes_value(true)
                .validator(can_be_parsed::<u8>),
        )
        .arg(
            Arg::with_name("separator")
                .short("s")
                .long("separator")
                .value_name("SEP")
                .default_value("-")
                .help("Separator between words")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("complexity")
                .short("c")
                .long("complexity")
                .value_name("COM")
                .possible_values(&["0", "1", "2"])
                .hide_possible_values(true)
                .default_value("0")
                .help("Use small words (0), medium words (1), or large words (2)")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("directory")
                .short("d")
                .long("dir")
                .value_name("DIR")
                .help("Directory containing adjectives.txt, adverbs.txt, names.txt")
                .conflicts_with("complexity")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("count")
                .long("count")
                .value_name("COUNT")
                .default_value("1")
                .help("Generate multiple names; pass 0 to produce infinite names!")
                .takes_value(true)
                .validator(can_be_parsed::<usize>),
        )
        .arg(
            Arg::with_name("letters")
                .short("l")
                .long("letters")
                .value_name("LETTERS")
                .default_value("0")
                .help("Maxiumum number of letters in each word; 0 for unlimited")
                .takes_value(true)
                .validator(can_be_parsed::<usize>),
        )
        .arg(
            Arg::with_name("alliterate")
                .short("a")
                .long("alliterate")
                .help("Generate names where each word begins with the same letter")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("alliterate-with")
                .long("alliterate-with")
                .short("A")
                .value_name("LETTER")
                .help("Generate names where each word begins with the given letter")
                .takes_value(true)
                .validator(can_be_parsed::<char>),
        )
        .arg(
            // For compatibility with upstream.
            Arg::with_name("ubuntu")
                .short("u")
                .long("ubuntu")
                .help("Alias; see --alliterate")
                .takes_value(false),
        )
}

fn run(matches: clap::ArgMatches) -> Result<(), Error> {
    // Unwrapping is safe because these options have defaults.
    let opt_separator = matches.value_of("separator").unwrap();
    let opt_words = matches.value_of("words").unwrap();
    let opt_complexity = matches.value_of("complexity").unwrap();
    let opt_count = matches.value_of("count").unwrap();
    let opt_letters = matches.value_of("letters").unwrap();
    let opt_alliterate = matches.is_present("alliterate")
        || matches.is_present("ubuntu")
        || matches.is_present("alliterate-with");

    // Optional arguments without defaults.
    let opt_directory = matches.value_of("directory");
    let opt_alliterate_char = matches
        .value_of("alliterate-with")
        .and_then(|s| s.parse::<char>().ok());

    // Parse numbers. Validated so unwrapping is okay.
    let opt_words: u8 = opt_words.parse().unwrap();
    let opt_count: usize = opt_count.parse().unwrap();
    let opt_letters: usize = opt_letters.parse().unwrap();

    // Load custom word lists, if specified.
    let words = match opt_directory {
        Some(dirname) => Words::load(dirname)?,
        None => Words::Builtin,
    };

    // Select the appropriate word list.
    let mut petnames = match words {
        Words::Custom(ref adjectives, ref adverbs, ref names) => {
            Petnames::init(adjectives, adverbs, names)
        }
        Words::Builtin => match opt_complexity {
            "0" => Petnames::small(),
            "1" => Petnames::medium(),
            "2" => Petnames::large(),
            _ => Petnames::small(),
        },
    };

    // If requested, limit the number of letters.
    if opt_letters != 0 {
        petnames.retain(|s| s.len() <= opt_letters);
    }

    // Check cardinality.
    if petnames.cardinality(opt_words) == 0 {
        return Err(Error::Cardinality(
            "no petnames to choose from; try relaxing constraints".to_string(),
        ));
    }

    // We're going to need a source of randomness.
    let mut rng = rand::thread_rng();

    // Handle alliteration, either by eliminating a specified
    // character, or using a random one.
    if opt_alliterate {
        // We choose the first letter from the intersection of the
        // first letters of each word list in `petnames`.
        let firsts =
            common_first_letters(&petnames.adjectives, &[&petnames.adverbs, &petnames.names]);
        // if a specific character was requested for alliteration,
        // attempt to use it.
        if let Some(c) = opt_alliterate_char {
            if firsts.contains(&c) {
                petnames.retain(|s| s.starts_with(c));
            } else {
                return Err(Error::Alliteration(
                    "no petnames begin with the choosen alliteration character".to_string(),
                ));
            }
        } else {
            // Otherwise choose the first letter at random; fails if
            // there are no letters.
            match firsts.iter().choose(&mut rng) {
                Some(c) => petnames.retain(|s| s.starts_with(*c)),
                None => {
                    return Err(Error::Alliteration(
                        "word lists have no initial letters in common".to_string(),
                    ))
                }
            };
        }
    }

    // Get an iterator for the names we want to print out.
    let names = petnames.iter(&mut rng, opt_words, opt_separator);

    // Manage stdout.
    let stdout = io::stdout();
    let mut writer = io::BufWriter::new(stdout.lock());

    // Stream if count is 0.
    if opt_count == 0 {
        for name in names {
            writeln!(writer, "{}", name).map_err(suppress_disconnect)?;
        }
    } else {
        for name in names.take(opt_count) {
            writeln!(writer, "{}", name)?;
        }
    }

    Ok(())
}

fn can_be_parsed<INTO>(value: String) -> Result<(), String>
where
    INTO: FromStr,
    <INTO as FromStr>::Err: std::fmt::Display,
{
    match value.parse::<INTO>() {
        Err(e) => Err(format!("{}", e)),
        Ok(_) => Ok(()),
    }
}

fn common_first_letters(init: &[&str], more: &[&[&str]]) -> HashSet<char> {
    let mut firsts = first_letters(init);
    let firsts_other: Vec<HashSet<char>> = more.iter().map(|list| first_letters(list)).collect();
    firsts.retain(|c| firsts_other.iter().all(|fs| fs.contains(c)));
    firsts
}

fn first_letters(names: &[&str]) -> HashSet<char> {
    names.iter().filter_map(|s| s.chars().next()).collect()
}

enum Words {
    Custom(String, String, String),
    Builtin,
}

impl Words {
    // Load word lists from the given directory. This function expects to find three
    // files in that directory: `adjectives.txt`, `adverbs.txt`, and `names.txt`.
    // Each should be valid UTF-8, and contain words separated by whitespace.
    fn load<T: AsRef<path::Path>>(dirname: T) -> Result<Self, Error> {
        let dirname = dirname.as_ref();
        Ok(Self::Custom(
            read_file_to_string(dirname.join("adjectives.txt"))?,
            read_file_to_string(dirname.join("adverbs.txt"))?,
            read_file_to_string(dirname.join("names.txt"))?,
        ))
    }
}

fn read_file_to_string<P: AsRef<path::Path>>(path: P) -> Result<String, Error> {
    fs::read_to_string(&path).map_err(|error| Error::FileIo(path.as_ref().to_path_buf(), error))
}

fn suppress_disconnect(err: io::Error) -> Error {
    match err.kind() {
        io::ErrorKind::BrokenPipe => Error::Disconnected,
        _ => err.into(),
    }
}
<commit_msg>Add --stream option; deprecate --count=0<commit_after>#[macro_use]
extern crate clap;

use petname::Petnames;

use std::collections::HashSet;
use std::fmt;
use std::fs;
use std::io::{self, Write};
use std::path;
use std::process;
use std::str::FromStr;

use clap::Arg;
use rand::seq::IteratorRandom;

fn main() {
    let matches = app().get_matches();
    match run(matches) {
        Err(Error::Disconnected) => {
            process::exit(0);
        }
        Err(e) => {
            eprintln!("Error: {}", e);
            process::exit(1);
        }
        Ok(()) => {
            process::exit(0);
        }
    }
}

enum Error {
    Io(io::Error),
    FileIo(path::PathBuf, io::Error),
    Cardinality(String),
    Alliteration(String),
    Disconnected,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Error::Io(ref e) => write!(f, "{}", e),
            Error::FileIo(ref path, ref e) => write!(f, "{}: {}", e, path.display()),
            Error::Cardinality(ref message) => write!(f, "cardinality is zero: {}", message),
            Error::Alliteration(ref message) => write!(f, "cannot alliterate: {}", message),
            Error::Disconnected => write!(f, "caller disconnected / stopped reading"),
        }
    }
}

impl From<io::Error> for Error {
    fn from(error: io::Error) -> Self {
        Error::Io(error)
    }
}

fn app<'a, 'b>() -> clap::App<'a, 'b> {
    clap::App::new("rust-petname")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Generate human readable random names.")
        .after_help(concat!(
            "Based on Dustin Kirkland's petname project ",
            "<https://github.com/dustinkirkland/petname>."
        ))
        .arg(
            Arg::with_name("words")
                .short("w")
                .long("words")
                .value_name("WORDS")
                .default_value("2")
                .help("Number of words in name")
                .takes_value(true)
                .validator(can_be_parsed::<u8>),
        )
        .arg(
            Arg::with_name("separator")
                .short("s")
                .long("separator")
                .value_name("SEP")
                .default_value("-")
                .help("Separator between words")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("complexity")
                .short("c")
                .long("complexity")
                .value_name("COM")
                .possible_values(&["0", "1", "2"])
                .hide_possible_values(true)
                .default_value("0")
                .help("Use small words (0), medium words (1), or large words (2)")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("directory")
                .short("d")
                .long("dir")
                .value_name("DIR")
                .help("Directory containing adjectives.txt, adverbs.txt, names.txt")
                .conflicts_with("complexity")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("count")
                .long("count")
                .value_name("COUNT")
                .default_value("1")
                .help(concat!(
                    "Generate multiple names; pass 0 to produce infinite ",
                    "names (--count=0 is deprecated; use --stream instead)"
                ))
                .takes_value(true)
                .validator(can_be_parsed::<usize>),
        )
        .arg(
            Arg::with_name("stream")
                .long("stream")
                .help("Stream names continuously")
                .conflicts_with("count")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("letters")
                .short("l")
                .long("letters")
                .value_name("LETTERS")
                .default_value("0")
                .help("Maxiumum number of letters in each word; 0 for unlimited")
                .takes_value(true)
                .validator(can_be_parsed::<usize>),
        )
        .arg(
            Arg::with_name("alliterate")
                .short("a")
                .long("alliterate")
                .help("Generate names where each word begins with the same letter")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("alliterate-with")
                .long("alliterate-with")
                .short("A")
                .value_name("LETTER")
                .help("Generate names where each word begins with the given letter")
                .takes_value(true)
                .validator(can_be_parsed::<char>),
        )
        .arg(
            // For compatibility with upstream.
            Arg::with_name("ubuntu")
                .short("u")
                .long("ubuntu")
                .help("Alias; see --alliterate")
                .takes_value(false),
        )
}

fn run(matches: clap::ArgMatches) -> Result<(), Error> {
    // Unwrapping is safe because these options have defaults.
    let opt_separator = matches.value_of("separator").unwrap();
    let opt_words = matches.value_of("words").unwrap();
    let opt_complexity = matches.value_of("complexity").unwrap();
    let opt_count = matches.value_of("count").unwrap();
    let opt_letters = matches.value_of("letters").unwrap();

    // Flags.
    let opt_stream = matches.is_present("stream");
    let opt_alliterate = matches.is_present("alliterate")
        || matches.is_present("ubuntu")
        || matches.is_present("alliterate-with");

    // Optional arguments without defaults.
    let opt_directory = matches.value_of("directory");
    let opt_alliterate_char = matches
        .value_of("alliterate-with")
        .and_then(|s| s.parse::<char>().ok());

    // Parse numbers. Validated so unwrapping is okay.
    let opt_words: u8 = opt_words.parse().unwrap();
    let opt_count: usize = opt_count.parse().unwrap();
    let opt_letters: usize = opt_letters.parse().unwrap();

    // Load custom word lists, if specified.
    let words = match opt_directory {
        Some(dirname) => Words::load(dirname)?,
        None => Words::Builtin,
    };

    // Select the appropriate word list.
    let mut petnames = match words {
        Words::Custom(ref adjectives, ref adverbs, ref names) => {
            Petnames::init(adjectives, adverbs, names)
        }
        Words::Builtin => match opt_complexity {
            "0" => Petnames::small(),
            "1" => Petnames::medium(),
            "2" => Petnames::large(),
            _ => Petnames::small(),
        },
    };

    // If requested, limit the number of letters.
    if opt_letters != 0 {
        petnames.retain(|s| s.len() <= opt_letters);
    }

    // Check cardinality.
    if petnames.cardinality(opt_words) == 0 {
        return Err(Error::Cardinality(
            "no petnames to choose from; try relaxing constraints".to_string(),
        ));
    }

    // We're going to need a source of randomness.
    let mut rng = rand::thread_rng();

    // Handle alliteration, either by eliminating a specified
    // character, or using a random one.
    if opt_alliterate {
        // We choose the first letter from the intersection of the
        // first letters of each word list in `petnames`.
        let firsts =
            common_first_letters(&petnames.adjectives, &[&petnames.adverbs, &petnames.names]);
        // if a specific character was requested for alliteration,
        // attempt to use it.
        if let Some(c) = opt_alliterate_char {
            if firsts.contains(&c) {
                petnames.retain(|s| s.starts_with(c));
            } else {
                return Err(Error::Alliteration(
                    "no petnames begin with the choosen alliteration character".to_string(),
                ));
            }
        } else {
            // Otherwise choose the first letter at random; fails if
            // there are no letters.
            match firsts.iter().choose(&mut rng) {
                Some(c) => petnames.retain(|s| s.starts_with(*c)),
                None => {
                    return Err(Error::Alliteration(
                        "word lists have no initial letters in common".to_string(),
                    ))
                }
            };
        }
    }

    // Get an iterator for the names we want to print out.
    let names = petnames.iter(&mut rng, opt_words, opt_separator);

    // Manage stdout.
    let stdout = io::stdout();
    let mut writer = io::BufWriter::new(stdout.lock());

    // Warn that --count=0 is deprecated.
    if opt_count == 0 {
        eprintln!(concat!(
            "Warning: specifying --count=0 to continuously produce petnames is ",
            "deprecated and its behaviour will change in a future version; ",
            "specify --stream instead.",
        ));
    }

    // Stream if count is 0. TODO: Only stream when --stream is specified.
    if opt_stream || opt_count == 0 {
        for name in names {
            writeln!(writer, "{}", name).map_err(suppress_disconnect)?;
        }
    } else {
        for name in names.take(opt_count) {
            writeln!(writer, "{}", name)?;
        }
    }

    Ok(())
}

fn can_be_parsed<INTO>(value: String) -> Result<(), String>
where
    INTO: FromStr,
    <INTO as FromStr>::Err: std::fmt::Display,
{
    match value.parse::<INTO>() {
        Err(e) => Err(format!("{}", e)),
        Ok(_) => Ok(()),
    }
}

fn common_first_letters(init: &[&str], more: &[&[&str]]) -> HashSet<char> {
    let mut firsts = first_letters(init);
    let firsts_other: Vec<HashSet<char>> = more.iter().map(|list| first_letters(list)).collect();
    firsts.retain(|c| firsts_other.iter().all(|fs| fs.contains(c)));
    firsts
}

fn first_letters(names: &[&str]) -> HashSet<char> {
    names.iter().filter_map(|s| s.chars().next()).collect()
}

enum Words {
    Custom(String, String, String),
    Builtin,
}

impl Words {
    // Load word lists from the given directory. This function expects to find three
    // files in that directory: `adjectives.txt`, `adverbs.txt`, and `names.txt`.
    // Each should be valid UTF-8, and contain words separated by whitespace.
    fn load<T: AsRef<path::Path>>(dirname: T) -> Result<Self, Error> {
        let dirname = dirname.as_ref();
        Ok(Self::Custom(
            read_file_to_string(dirname.join("adjectives.txt"))?,
            read_file_to_string(dirname.join("adverbs.txt"))?,
            read_file_to_string(dirname.join("names.txt"))?,
        ))
    }
}

fn read_file_to_string<P: AsRef<path::Path>>(path: P) -> Result<String, Error> {
    fs::read_to_string(&path).map_err(|error| Error::FileIo(path.as_ref().to_path_buf(), error))
}

fn suppress_disconnect(err: io::Error) -> Error {
    match err.kind() {
        io::ErrorKind::BrokenPipe => Error::Disconnected,
        _ => err.into(),
    }
}
<|endoftext|>


<commit_before>//! An implementation of [tldr](https://github.com/tldr-pages/tldr) in Rust.

#[macro_use] extern crate log;
#[cfg(feature = "logging")]extern crate env_logger;
extern crate ansi_term;

use std::io::{BufRead, BufReader};
use std::fs::File;
use std::{env, process};

mod types;
mod tokenizer;

use types::LineType;
use tokenizer::Tokenizer;

/// Open file, return a `BufRead` instance
fn get_file_reader(filepath: &str) -> Result<BufReader<File>, String> {
    let file = try!(
        File::open(filepath)
            .map_err(|msg| format!("Could not open file: {}", msg))
    );
    Ok(BufReader::new(file))
}

/// Print a token stream to an ANSI terminal.
fn print_lines<R>(tokenizer: &mut Tokenizer<R>) where R: BufRead {
    while let Some(token) = tokenizer.next() {
        match token {
            LineType::Empty => println!(""),
            LineType::Title(_) => debug!("Ignoring title"),
            LineType::Description(text) => println!("  {}", text),
            LineType::ExampleText(text) => println!("  - {}", text),
            LineType::ExampleCode(text) => println!("    {}", text),
            LineType::Other(text) => debug!("Unknown line type: {:?}", text),
        }
    }
}

#[cfg(feature = "logging")]
fn init_log() {
    env_logger::init().unwrap();
}

#[cfg(not(feature = "logging"))]
fn init_log() { }

fn main() {

    // Initialize logger
    init_log();

    // Parse arguments
    let args: Vec<_> = env::args().collect();
    if args.len() != 2 {
        println!("Usage: {} <command>", args[0]);
        process::exit(1);
    }

    // Open file
    let reader = get_file_reader(&args[1]).unwrap_or_else(|msg| {
        println!("{}", msg);
        process::exit(1);
    });

    // Create tokenizer
    let mut tokenizer = Tokenizer::new(reader);

    // Print output
    print_lines(&mut tokenizer);

}

#[cfg(test)]
mod test {
    use types::LineType;

    #[test]
    fn test_linetype_from_str() {
        assert_eq!(LineType::from(""), LineType::Empty);
        assert_eq!(LineType::from(" \n \r"), LineType::Empty);
        assert_eq!(LineType::from("# Hello there"), LineType::Title("Hello there".into()));
        assert_eq!(LineType::from("> tis a description \n"), LineType::Description("tis a description".into()));
        assert_eq!(LineType::from("- some command"), LineType::ExampleText("some command".into()));
        assert_eq!(LineType::from("`$ cargo run`"), LineType::ExampleCode("$ cargo run".into()));
        assert_eq!(LineType::from("`$ cargo run"), LineType::Other("`$ cargo run".into()));
        assert_eq!(LineType::from("jklö"), LineType::Other("jklö".into()));
    }
}
<commit_msg>Add color support<commit_after>//! An implementation of [tldr](https://github.com/tldr-pages/tldr) in Rust.

#[macro_use] extern crate log;
#[cfg(feature = "logging")]extern crate env_logger;
extern crate ansi_term;

use std::io::{BufRead, BufReader};
use std::fs::File;
use std::{env, process};

use ansi_term::Colour;

mod types;
mod tokenizer;

use types::LineType;
use tokenizer::Tokenizer;

/// Open file, return a `BufRead` instance
fn get_file_reader(filepath: &str) -> Result<BufReader<File>, String> {
    let file = try!(
        File::open(filepath)
            .map_err(|msg| format!("Could not open file: {}", msg))
    );
    Ok(BufReader::new(file))
}

/// Print a token stream to an ANSI terminal.
fn print_lines<R>(tokenizer: &mut Tokenizer<R>) where R: BufRead {
    while let Some(token) = tokenizer.next() {
        match token {
            LineType::Empty => println!(""),
            LineType::Title(_) => debug!("Ignoring title"),
            LineType::Description(text) => println!("  {}", text),
            LineType::ExampleText(text) => println!("  {}", Colour::Green.paint(format!("- {}", text))),
            LineType::ExampleCode(text) => println!("  {}", Colour::Cyan.paint(format!("  {}", text))),
            LineType::Other(text) => debug!("Unknown line type: {:?}", text),
        }
    }
}

#[cfg(feature = "logging")]
fn init_log() {
    env_logger::init().unwrap();
}

#[cfg(not(feature = "logging"))]
fn init_log() { }

fn main() {

    // Initialize logger
    init_log();

    // Parse arguments
    let args: Vec<_> = env::args().collect();
    if args.len() != 2 {
        println!("Usage: {} <command>", args[0]);
        process::exit(1);
    }

    // Open file
    let reader = get_file_reader(&args[1]).unwrap_or_else(|msg| {
        println!("{}", msg);
        process::exit(1);
    });

    // Create tokenizer
    let mut tokenizer = Tokenizer::new(reader);

    // Print output
    print_lines(&mut tokenizer);
    println!("");

}

#[cfg(test)]
mod test {
    use types::LineType;

    #[test]
    fn test_linetype_from_str() {
        assert_eq!(LineType::from(""), LineType::Empty);
        assert_eq!(LineType::from(" \n \r"), LineType::Empty);
        assert_eq!(LineType::from("# Hello there"), LineType::Title("Hello there".into()));
        assert_eq!(LineType::from("> tis a description \n"), LineType::Description("tis a description".into()));
        assert_eq!(LineType::from("- some command"), LineType::ExampleText("some command".into()));
        assert_eq!(LineType::from("`$ cargo run`"), LineType::ExampleCode("$ cargo run".into()));
        assert_eq!(LineType::from("`$ cargo run"), LineType::Other("`$ cargo run".into()));
        assert_eq!(LineType::from("jklö"), LineType::Other("jklö".into()));
    }
}
<|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 = regex::Regex::new(&ctx.get::<Box<str>>(0)?)
                .map_err(|e| rusqlite::Error::UserFunctionError(Box::new(e)))?;
            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; 4] = [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>Cache teh regexp for speed. (#549)<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.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; 4] = [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|>