<commit_before>#![feature(plugin)]
#![plugin(regex_macros)]

extern crate argparse;
extern crate html5ever;
extern crate html5ever_dom_sink;
extern crate hyper;
extern crate iron;
#[macro_use]
extern crate log;
extern crate log4rs;
extern crate rand;
extern crate regex;
extern crate rustc_serialize;
extern crate tempfile;
extern crate tendril;
extern crate url;

/// Helper macro for unwrapping Result values whose E types implement std::fmt::Display. For Ok(),
/// evaluates to the contained value. For Err(), returns early with an Err containing the formatted
/// error.
macro_rules! try_display {
    ($expr:expr, $($format_arg:expr),* ) => (match $expr {
        Ok(val) => val,
        Err(err) => return Err(format!("{}: {}", format!($($format_arg),*), err)),
    })
}

mod json;
mod wiki;

use argparse::ArgumentParser;
use argparse::Store;
use std::fs::OpenOptions;
use std::io::Read;
use std::io::Write;
use std::process::Command;
use std::process::Stdio;
use std::str::FromStr;

use html5ever::Attribute;
use html5ever::tree_builder::interface::TreeSink;
use html5ever_dom_sink::common::NodeEnum;
use html5ever_dom_sink::rcdom::Handle;
use html5ever_dom_sink::rcdom::RcDom;
use hyper::Client;
use hyper::header::Connection;
use iron::Iron;
use iron::IronResult;
use iron::Request;
use iron::Response;
use iron::headers::ContentType;
use iron::middleware::Handler;
use iron::mime::Mime;
use iron::mime::SubLevel;
use iron::mime::TopLevel;
use tempfile::NamedTempFile;

use wiki::Revision;
use wiki::Wiki;

// TODO: consider doing s/en.wikipedia.org/this app's url/ on the HTML before serving it. This
// currently works fine, but might not over HTTPS.

/// Does a 3-way merge, merging `new1` and `new2` under the assumption that both diverged from
/// `old`. Returns None if the strings do not merge together cleanly.
fn merge(old: &str, new1: &str, new2: &str) -> Result<Option<String>, String> {
    fn write_to_temp_file(contents: &str) -> Result<NamedTempFile, String> {
        let tempfile = try_display!(NamedTempFile::new(), "Error creating temp file");
        let mut file = try_display!(OpenOptions::new().write(true).open(tempfile.path()),
                                    "Error opening temp file for writing");
        try_display!(file.write_all(contents.as_bytes()), "Error writing to temp file");
        try_display!(file.flush(), "Error flushing temp file");
        Ok(tempfile)
    }

    let old_tempfile = try!(write_to_temp_file(old));
    let new1_tempfile = try!(write_to_temp_file(new1));
    let new2_tempfile = try!(write_to_temp_file(new2));

    let mut process = Command::new("diff3");
    process.arg("-m").args(&[new1_tempfile.path(), old_tempfile.path(), new2_tempfile.path()])
        .stdout(Stdio::piped()).stderr(Stdio::null());
    let output = try_display!(process.output(), "Error getting output from diff3 subprocess");
    if output.status.success() {
        Ok(Some(try_display!(String::from_utf8(output.stdout),
                             "Error converting diff3 output to UTF-8")))
    } else {
        Ok(None)
    }
}

fn replace_node_with_placeholder(original_html: &str, div_id: &str, placeholder: &str)
    -> Result<String, String> {
    let html = tendril::StrTendril::from_str(original_html).unwrap();
    let mut dom: RcDom = html5ever::parse(html5ever::one_input(html), Default::default());

    let handle = try!(find_node_by_id(&dom.get_document(), div_id));
    let child_handles =
        (&handle.borrow().children).into_iter().map(|child| child.clone()).collect::<Vec<_>>();
    for child_handle in child_handles {
        dom.remove_from_parent(child_handle);
    }
    dom.append(handle,
               html5ever::tree_builder::interface::NodeOrText::AppendText(
                   tendril::StrTendril::from_str(placeholder).unwrap()));
    let mut serialized: Vec<u8> = vec![];
    try_display!(
        html5ever::serialize::serialize(&mut serialized, &dom.document, Default::default()),
        "Failed to serialize modified HTML");
    Ok(try_display!(String::from_utf8(serialized),
                    "Error converting serialized HTML to UTF-8 string"))
}

fn find_node_by_id(handle: &Handle, id: &str) -> Result<Handle, String> {
    fn has_matching_id(attributes: &Vec<Attribute>, id: &str) -> bool {
        return attributes.into_iter().any(
            |attribute| attribute.name.local.as_slice() == "id" &&
                format!("{}", attribute.value) == id);
    }

    let node = handle.borrow();
    match node.node {
        NodeEnum::Element(_, ref attributes) if has_matching_id(attributes, id) => Ok(handle.clone()),
        _ => (&node.children).into_iter()
            .map(|child| find_node_by_id(child, id))
            .filter(|result| result.is_ok())
            .map(|result| result.unwrap())
            .next().ok_or(format!("No node with ID {} found", id)),
    }
}

struct WikipediaMinusWikipediansHandler {
    wiki: Wiki,
    client: Client
}

impl WikipediaMinusWikipediansHandler {
    fn new(wiki: Wiki, client: Client) -> WikipediaMinusWikipediansHandler {
        WikipediaMinusWikipediansHandler { wiki: wiki, client: client }
    }

    /// Returns a vector of Revisions representing all reversions of vandalism for the page `title`.
    fn get_vandalism_reversions(&self, title: &str) -> Result<Vec<Revision>, String> {
        // TODO: Add a flag to control number of revisions, and experiment with it.
        let revisions = try!(self.wiki.get_revisions(title, 60));
        Ok(revisions.into_iter().filter(|revision| revision.comment.contains("vandal")).collect())
    }

    // TODO: this name is terrible.
    fn get_page_with_vandalism_restored(&self, title: &str) -> Result<String, String> {
        let canonical_title = try!(self.wiki.get_canonical_title(title));
        info!("Canonical page title for \"{}\" is \"{}\"", title, canonical_title);

        let latest_revid = try!(self.wiki.get_latest_revision(&canonical_title)).revid;
        let mut accumulated_contents =
            try!(self.wiki.get_revision_content(&canonical_title, latest_revid));
        let mut merged_revisions = vec![];
        let vandalism_reversions = try!(self.get_vandalism_reversions(&canonical_title));
        for revision in &vandalism_reversions {
            let reverted_contents =
                try!(self.wiki.get_revision_content(&canonical_title, revision.revid));
            let vandalized_contents =
                try!(self.wiki.get_revision_content(&canonical_title, revision.parentid));
            match merge(&reverted_contents, &vandalized_contents, &accumulated_contents) {
                Ok(Some(merged_contents)) => {
                    accumulated_contents = merged_contents;
                    merged_revisions.push(revision.revid);
                }
                Ok(None) => (),
                Err(msg) => return Err(format!("Error merging revision {} of \"{}\": {}",
                                               revision.revid, title, msg))
            }
        }

        info!("For \"{}\", successfully merged {} of {} possible vandalisms: {:?}", title,
              merged_revisions.len(), vandalism_reversions.len(), merged_revisions);

        let body = try!(self.wiki.parse_wikitext(&canonical_title, &accumulated_contents));
        // Note: "title" rather than "canonical_title", so that redirects look right.
        let current_page_contents = try!(self.wiki.get_current_page_content(&title));
        let placeholder = format!("WMW_PLACEHOLDER_TEXT_{}", rand::random::<u64>());
        let page_contents_with_placeholder =
            try!(replace_node_with_placeholder(
                &current_page_contents, "mw-content-text", &placeholder));
        Ok(page_contents_with_placeholder.replace(&placeholder, &body))
    }
}

impl Handler for WikipediaMinusWikipediansHandler {
    fn handle(&self, request: &mut Request) -> IronResult<Response> {
        if request.url.path.len() == 2 && request.url.path[0] == "wiki" {
            let mut response =
                match self.get_page_with_vandalism_restored(&request.url.path[1]) {
                    Ok(page_contents) => Response::with((iron::status::Ok, page_contents)),
                    // TODO: create an Error type to pass around, so this can distinguish different
                    // types of error (if that would be helpful).
                    // TODO: create a better error page
                    Err(msg) => {
                        warn!("Failed to get page with vandalism restored: {}", msg);
                        Response::with(
                            (iron::status::InternalServerError, "ERROR"))
                    },
                };
            response.headers.set(ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![])));
            Ok(response)
        } else {
            // TODO: should I use an HTTP redirect here instead? Would that work? Would it be desirable?
            // TODO: Maybe should be moved to wiki module.
            let mut url = request.url.clone();
            url.scheme = "https".to_string();
            url.host = url::Host::Domain(self.wiki.hostname.clone());
            url.port = self.wiki.port;
            let url = url.into_generic_url().serialize();
            match self.client.get(&url)
                .header(Connection::close()).send() {
                    Ok(mut wikipedia_response) => {
                        let mut wikipedia_body: Vec<u8> = Vec::new();
                        match wikipedia_response.read_to_end(&mut wikipedia_body) {
                            Ok(..) => {
                                info!("Received {} response from {}", wikipedia_response.status,
                                      url);
                                let mut response = Response::with(wikipedia_body);
                                response.status = Some(wikipedia_response.status);
                                response.headers = wikipedia_response.headers.clone();
                                Ok(response)
                            },
                            Err(error) => {
                                warn!("Error reading Wikipedia response: {}", error);
                                let mut response = Response::with(
                                    (iron::status::InternalServerError,
                                     "ERROR"));
                                response.headers.set(
                                    ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![])));
                                Ok(response)
                            }
                        }
                    },
                    Err(error) => {
                        warn!("Error reading URL {}: {}", url, error);
                        let mut response = Response::with(
                            (iron::status::InternalServerError,
                             "ERROR: {}"));
                        response.headers.set(
                            ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![])));
                        Ok(response)
                    }
                }
        }
    }
}

fn main() {
    log4rs::init_file("log.toml", Default::default()).unwrap();

    let mut port = 3000;
    let mut wiki = "en.wikipedia.org".to_string();
    {
        let mut parser = ArgumentParser::new();
        parser.set_description("TODO: Usage description");
        parser.refer(&mut port).add_option(&["-p", "--port"], Store, "The port to serve HTTP on.");
        parser.refer(&mut wiki).add_option(
            &["--wiki"], Store, "The hostname or hostname:port of the wiki to mirror.");
        parser.parse_args_or_exit();
    }
    let mut wiki_components = wiki.split(":");
    let wiki_hostname = wiki_components.next().unwrap();
    let wiki_port = match wiki_components.next() {
        Some(port) => port.parse::<u16>().unwrap(),
        None => 443,
    };

    let handler =
        WikipediaMinusWikipediansHandler::new(
            Wiki::new(wiki_hostname.to_string(), wiki_port, Client::new()), Client::new());
    Iron::new(handler).http(("localhost", port)).unwrap();
}

#[cfg(test)]
mod tests {
    use super::{merge, replace_node_with_placeholder};

    #[test]
    fn test_merge_clean() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed.\n";
        let new2 = "First line changed.\n\nSecond line.\n";
        assert_eq!(Some("First line changed.\n\nSecond line changed.\n".to_string()),
                   merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_merge_conflicting() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed one way.\n";
        let new2 = "First line changed.\n\nSecond line changed a different way.\n";
        assert_eq!(None, merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_merge_special_characters() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line 𐅃.\n";
        let new2 = "First line さようなら.\n\nSecond line.\n";
        assert_eq!(Some("First line さようなら.\n\nSecond line 𐅃.\n".to_string()),
                   merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_replace_html_content() {
        let original_html = "

original text

Other text

";
        let expected_html = "

replaced text

Other text

";
        let processed_html = replace_node_with_placeholder(original_html, "mw-content-text", "replaced text").unwrap();
        assert_eq!(expected_html, processed_html);
    }
}
<commit_msg>Use 500 (the limit on Wikipedia) for the revision limit in Wiki::get_revisions().<commit_after>#![feature(plugin)]
#![plugin(regex_macros)]

extern crate argparse;
extern crate html5ever;
extern crate html5ever_dom_sink;
extern crate hyper;
extern crate iron;
#[macro_use]
extern crate log;
extern crate log4rs;
extern crate rand;
extern crate regex;
extern crate rustc_serialize;
extern crate tempfile;
extern crate tendril;
extern crate url;

/// Helper macro for unwrapping Result values whose E types implement std::fmt::Display. For Ok(),
/// evaluates to the contained value. For Err(), returns early with an Err containing the formatted
/// error.
macro_rules! try_display {
    ($expr:expr, $($format_arg:expr),* ) => (match $expr {
        Ok(val) => val,
        Err(err) => return Err(format!("{}: {}", format!($($format_arg),*), err)),
    })
}

mod json;
mod wiki;

use argparse::ArgumentParser;
use argparse::Store;
use std::fs::OpenOptions;
use std::io::Read;
use std::io::Write;
use std::process::Command;
use std::process::Stdio;
use std::str::FromStr;

use html5ever::Attribute;
use html5ever::tree_builder::interface::TreeSink;
use html5ever_dom_sink::common::NodeEnum;
use html5ever_dom_sink::rcdom::Handle;
use html5ever_dom_sink::rcdom::RcDom;
use hyper::Client;
use hyper::header::Connection;
use iron::Iron;
use iron::IronResult;
use iron::Request;
use iron::Response;
use iron::headers::ContentType;
use iron::middleware::Handler;
use iron::mime::Mime;
use iron::mime::SubLevel;
use iron::mime::TopLevel;
use tempfile::NamedTempFile;

use wiki::Revision;
use wiki::Wiki;

// TODO: consider doing s/en.wikipedia.org/this app's url/ on the HTML before serving it. This
// currently works fine, but might not over HTTPS.

/// Does a 3-way merge, merging `new1` and `new2` under the assumption that both diverged from
/// `old`. Returns None if the strings do not merge together cleanly.
fn merge(old: &str, new1: &str, new2: &str) -> Result<Option<String>, String> {
    fn write_to_temp_file(contents: &str) -> Result<NamedTempFile, String> {
        let tempfile = try_display!(NamedTempFile::new(), "Error creating temp file");
        let mut file = try_display!(OpenOptions::new().write(true).open(tempfile.path()),
                                    "Error opening temp file for writing");
        try_display!(file.write_all(contents.as_bytes()), "Error writing to temp file");
        try_display!(file.flush(), "Error flushing temp file");
        Ok(tempfile)
    }

    let old_tempfile = try!(write_to_temp_file(old));
    let new1_tempfile = try!(write_to_temp_file(new1));
    let new2_tempfile = try!(write_to_temp_file(new2));

    let mut process = Command::new("diff3");
    process.arg("-m").args(&[new1_tempfile.path(), old_tempfile.path(), new2_tempfile.path()])
        .stdout(Stdio::piped()).stderr(Stdio::null());
    let output = try_display!(process.output(), "Error getting output from diff3 subprocess");
    if output.status.success() {
        Ok(Some(try_display!(String::from_utf8(output.stdout),
                             "Error converting diff3 output to UTF-8")))
    } else {
        Ok(None)
    }
}

fn replace_node_with_placeholder(original_html: &str, div_id: &str, placeholder: &str)
    -> Result<String, String> {
    let html = tendril::StrTendril::from_str(original_html).unwrap();
    let mut dom: RcDom = html5ever::parse(html5ever::one_input(html), Default::default());

    let handle = try!(find_node_by_id(&dom.get_document(), div_id));
    let child_handles =
        (&handle.borrow().children).into_iter().map(|child| child.clone()).collect::<Vec<_>>();
    for child_handle in child_handles {
        dom.remove_from_parent(child_handle);
    }
    dom.append(handle,
               html5ever::tree_builder::interface::NodeOrText::AppendText(
                   tendril::StrTendril::from_str(placeholder).unwrap()));
    let mut serialized: Vec<u8> = vec![];
    try_display!(
        html5ever::serialize::serialize(&mut serialized, &dom.document, Default::default()),
        "Failed to serialize modified HTML");
    Ok(try_display!(String::from_utf8(serialized),
                    "Error converting serialized HTML to UTF-8 string"))
}

fn find_node_by_id(handle: &Handle, id: &str) -> Result<Handle, String> {
    fn has_matching_id(attributes: &Vec<Attribute>, id: &str) -> bool {
        return attributes.into_iter().any(
            |attribute| attribute.name.local.as_slice() == "id" &&
                format!("{}", attribute.value) == id);
    }

    let node = handle.borrow();
    match node.node {
        NodeEnum::Element(_, ref attributes) if has_matching_id(attributes, id) => Ok(handle.clone()),
        _ => (&node.children).into_iter()
            .map(|child| find_node_by_id(child, id))
            .filter(|result| result.is_ok())
            .map(|result| result.unwrap())
            .next().ok_or(format!("No node with ID {} found", id)),
    }
}

struct WikipediaMinusWikipediansHandler {
    wiki: Wiki,
    client: Client
}

impl WikipediaMinusWikipediansHandler {
    fn new(wiki: Wiki, client: Client) -> WikipediaMinusWikipediansHandler {
        WikipediaMinusWikipediansHandler { wiki: wiki, client: client }
    }

    /// Returns a vector of Revisions representing all reversions of vandalism for the page `title`.
    fn get_vandalism_reversions(&self, title: &str) -> Result<Vec<Revision>, String> {
        let revisions = try!(self.wiki.get_revisions(title, 500));
        Ok(revisions.into_iter().filter(|revision| revision.comment.contains("vandal")).collect())
    }

    // TODO: this name is terrible.
    fn get_page_with_vandalism_restored(&self, title: &str) -> Result<String, String> {
        let canonical_title = try!(self.wiki.get_canonical_title(title));
        info!("Canonical page title for \"{}\" is \"{}\"", title, canonical_title);

        let latest_revid = try!(self.wiki.get_latest_revision(&canonical_title)).revid;
        let mut accumulated_contents =
            try!(self.wiki.get_revision_content(&canonical_title, latest_revid));
        let mut merged_revisions = vec![];
        let vandalism_reversions = try!(self.get_vandalism_reversions(&canonical_title));
        for revision in &vandalism_reversions {
            let reverted_contents =
                try!(self.wiki.get_revision_content(&canonical_title, revision.revid));
            let vandalized_contents =
                try!(self.wiki.get_revision_content(&canonical_title, revision.parentid));
            match merge(&reverted_contents, &vandalized_contents, &accumulated_contents) {
                Ok(Some(merged_contents)) => {
                    accumulated_contents = merged_contents;
                    merged_revisions.push(revision.revid);
                }
                Ok(None) => (),
                Err(msg) => return Err(format!("Error merging revision {} of \"{}\": {}",
                                               revision.revid, title, msg))
            }
        }

        info!("For \"{}\", successfully merged {} of {} possible vandalisms: {:?}", title,
              merged_revisions.len(), vandalism_reversions.len(), merged_revisions);

        let body = try!(self.wiki.parse_wikitext(&canonical_title, &accumulated_contents));
        // Note: "title" rather than "canonical_title", so that redirects look right.
        let current_page_contents = try!(self.wiki.get_current_page_content(&title));
        let placeholder = format!("WMW_PLACEHOLDER_TEXT_{}", rand::random::<u64>());
        let page_contents_with_placeholder =
            try!(replace_node_with_placeholder(
                &current_page_contents, "mw-content-text", &placeholder));
        Ok(page_contents_with_placeholder.replace(&placeholder, &body))
    }
}

impl Handler for WikipediaMinusWikipediansHandler {
    fn handle(&self, request: &mut Request) -> IronResult<Response> {
        if request.url.path.len() == 2 && request.url.path[0] == "wiki" {
            let mut response =
                match self.get_page_with_vandalism_restored(&request.url.path[1]) {
                    Ok(page_contents) => Response::with((iron::status::Ok, page_contents)),
                    // TODO: create an Error type to pass around, so this can distinguish different
                    // types of error (if that would be helpful).
                    // TODO: create a better error page
                    Err(msg) => {
                        warn!("Failed to get page with vandalism restored: {}", msg);
                        Response::with(
                            (iron::status::InternalServerError, "ERROR"))
                    },
                };
            response.headers.set(ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![])));
            Ok(response)
        } else {
            // TODO: should I use an HTTP redirect here instead? Would that work? Would it be desirable?
            // TODO: Maybe should be moved to wiki module.
            let mut url = request.url.clone();
            url.scheme = "https".to_string();
            url.host = url::Host::Domain(self.wiki.hostname.clone());
            url.port = self.wiki.port;
            let url = url.into_generic_url().serialize();
            match self.client.get(&url)
                .header(Connection::close()).send() {
                    Ok(mut wikipedia_response) => {
                        let mut wikipedia_body: Vec<u8> = Vec::new();
                        match wikipedia_response.read_to_end(&mut wikipedia_body) {
                            Ok(..) => {
                                info!("Received {} response from {}", wikipedia_response.status,
                                      url);
                                let mut response = Response::with(wikipedia_body);
                                response.status = Some(wikipedia_response.status);
                                response.headers = wikipedia_response.headers.clone();
                                Ok(response)
                            },
                            Err(error) => {
                                warn!("Error reading Wikipedia response: {}", error);
                                let mut response = Response::with(
                                    (iron::status::InternalServerError,
                                     "ERROR"));
                                response.headers.set(
                                    ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![])));
                                Ok(response)
                            }
                        }
                    },
                    Err(error) => {
                        warn!("Error reading URL {}: {}", url, error);
                        let mut response = Response::with(
                            (iron::status::InternalServerError,
                             "ERROR: {}"));
                        response.headers.set(
                            ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![])));
                        Ok(response)
                    }
                }
        }
    }
}

fn main() {
    log4rs::init_file("log.toml", Default::default()).unwrap();

    let mut port = 3000;
    let mut wiki = "en.wikipedia.org".to_string();
    {
        let mut parser = ArgumentParser::new();
        parser.set_description("TODO: Usage description");
        parser.refer(&mut port).add_option(&["-p", "--port"], Store, "The port to serve HTTP on.");
        parser.refer(&mut wiki).add_option(
            &["--wiki"], Store, "The hostname or hostname:port of the wiki to mirror.");
        parser.parse_args_or_exit();
    }
    let mut wiki_components = wiki.split(":");
    let wiki_hostname = wiki_components.next().unwrap();
    let wiki_port = match wiki_components.next() {
        Some(port) => port.parse::<u16>().unwrap(),
        None => 443,
    };

    let handler =
        WikipediaMinusWikipediansHandler::new(
            Wiki::new(wiki_hostname.to_string(), wiki_port, Client::new()), Client::new());
    Iron::new(handler).http(("localhost", port)).unwrap();
}

#[cfg(test)]
mod tests {
    use super::{merge, replace_node_with_placeholder};

    #[test]
    fn test_merge_clean() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed.\n";
        let new2 = "First line changed.\n\nSecond line.\n";
        assert_eq!(Some("First line changed.\n\nSecond line changed.\n".to_string()),
                   merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_merge_conflicting() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed one way.\n";
        let new2 = "First line changed.\n\nSecond line changed a different way.\n";
        assert_eq!(None, merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_merge_special_characters() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line 𐅃.\n";
        let new2 = "First line さようなら.\n\nSecond line.\n";
        assert_eq!(Some("First line さようなら.\n\nSecond line 𐅃.\n".to_string()),
                   merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_replace_html_content() {
        let original_html = "

original text

Other text

";
        let expected_html = "

replaced text

Other text

";
        let processed_html = replace_node_with_placeholder(original_html, "mw-content-text", "replaced text").unwrap();
        assert_eq!(expected_html, processed_html);
    }
}
<|endoftext|>


<commit_before>//! Documentation for Rust Bio Tools
use clap::{load_yaml, value_t};
use log::LevelFilter;

use clap::App;
use itertools::Itertools;
use std::collections::HashMap;
use std::error::Error;
use std::fs;
use std::path::Path;
use std::str::FromStr;

pub mod bam;
pub mod bcf;
pub mod common;
pub mod fastq;
pub mod sequences_stats;

fn main() -> Result<(), Box<dyn Error>> {
    let yaml = load_yaml!("cli.yaml");
    let matches = App::from_yaml(yaml)
        .version(env!("CARGO_PKG_VERSION"))
        .get_matches();

    fern::Dispatch::new()
        .format(|out, message, _| out.finish(format_args!("{}", message)))
        .level(if matches.is_present("verbose") {
            LevelFilter::Debug
        } else {
            LevelFilter::Info
        })
        .chain(std::io::stderr())
        .apply()
        .unwrap();

    match matches.subcommand() {
        ("fastq-split", Some(matches)) => {
            fastq::split::split(&matches.values_of("chunks").unwrap().collect_vec())
        }
        ("fastq-filter", Some(matches)) => fastq::filter::filter(&matches.value_of("ids").unwrap()),
        ("bam-depth", Some(matches)) => bam::depth::depth(
            &matches.value_of("bam-path").unwrap(),
            value_t!(matches, "max-read-length", u32).unwrap_or(1000),
            value_t!(matches, "include-flags", u16).unwrap_or(0),
            value_t!(matches, "exclude-flags", u16).unwrap_or(4 | 256 | 512 | 1024),
            value_t!(matches, "min-mapq", u8).unwrap_or(0),
        ),
        ("vcf-to-txt", Some(matches)) => bcf::to_txt::to_txt(
            &matches
                .values_of("info")
                .map(|values| values.collect_vec())
                .unwrap_or_default(),
            &matches
                .values_of("format")
                .map(|values| values.collect_vec())
                .unwrap_or_default(),
            matches.is_present("genotypes"),
        ),
        ("vcf-match", Some(matches)) => bcf::match_variants::match_variants(
            matches.value_of("vcf").unwrap(),
            value_t!(matches, "max-dist", u32).unwrap_or(20),
            value_t!(matches, "max-len-diff", u32).unwrap_or(10),
        ),
        ("vcf-baf", Some(_)) => bcf::baf::calculate_baf(),
        ("vcf-fix-iupac-alleles", Some(_)) => bcf::fix_iupac_alleles::fix_iupac_alleles(),
        ("vcf-annotate-dgidb", Some(matches)) => bcf::annotate_dgidb::annotate_dgidb(
            &matches.value_of("vcf").unwrap(),
            matches.value_of("api-path").unwrap().to_string(),
            &matches.value_of("field").unwrap(),
            match matches.is_present("datasources") {
                true => Some(matches.values_of("datasources").unwrap().collect()),
                false => None,
            },
            value_t!(matches, "genes-per-request", usize).unwrap(),
        ),
        ("vcf-report", Some(matches)) => {
            let mut sample_calls = HashMap::new();
            let mut bam_paths = HashMap::new();
            let output_path = matches.value_of("output-path").unwrap();
            let max_cells = u32::from_str(matches.value_of("max-cells").unwrap()).unwrap();
            let custom_js = matches.value_of("custom-js");
            bcf::report::embed_js(output_path, custom_js)?;
            bcf::report::embed_css(output_path)?;
            bcf::report::embed_html(output_path)?;
            let fasta_path = matches.value_of("fasta").unwrap();
            let detail_path = output_path.to_owned() + "/details/";
            fs::create_dir(Path::new(&detail_path))?;
            let vcfs = matches.values_of("vcfs").unwrap();
            let bams = matches.values_of("bams").unwrap();
            let infos = matches.values_of("info");
            let formats = matches.values_of("format");
            for vcf in vcfs {
                let v: Vec<_> = vcf.split('=').collect();
                match sample_calls.insert(v[0].to_owned(), v[1].to_owned()) {
                    None => {}
                    _ => panic!("Found duplicate sample name {}. Please make sure the provided sample names are unique.", v[0].to_owned())
                }
            }
            for bam in bams {
                let b: Vec<_> = bam.split('=').collect();
                let c: Vec<_> = b[0].split(':').collect();
                let rec = bam_paths.entry(c[0].to_owned()).or_insert_with(Vec::new);
                rec.push((c[1].to_owned(), b[1].to_owned()))
            }

            for sample in sample_calls.keys().sorted() {
                bcf::report::table_report::table_report(
                    sample_calls.get(sample).unwrap(),
                    fasta_path,
                    bam_paths
                        .get(sample)
                        .unwrap_or_else(|| panic!("No bam provided for sample {}", sample)),
                    output_path,
                    sample,
                    infos.clone(),
                    formats.clone(),
                )?;
            }

            bcf::report::oncoprint::oncoprint(&sample_calls, output_path, max_cells)
        }
        ("collapse-reads-to-fragments", Some(matches)) => match matches.subcommand() {
            ("fastq", Some(matches)) => {
                fastq::collapse_reads_to_fragments::call_consensus_reads_from_paths(
                    matches.value_of("fq1").unwrap(),
                    matches.value_of("fq2").unwrap(),
                    matches.value_of("consensus-fq1").unwrap(),
                    matches.value_of("consensus-fq2").unwrap(),
                    matches.value_of("consensus-fq3"),
                    value_t!(matches, "umi-len", usize).unwrap(),
                    value_t!(matches, "max-seq-dist", usize).unwrap(),
                    value_t!(matches, "max-umi-dist", usize).unwrap(),
                    matches.is_present("umi-on-reverse"),
                    matches.is_present("verbose-read-names"),
                    if matches.is_present("insert-size") {
                        Some(value_t!(matches, "insert-size", usize).unwrap())
                    } else {
                        None
                    },
                    if matches.is_present("std-dev") {
                        Some(value_t!(matches, "std-dev", usize).unwrap())
                    } else {
                        None
                    },
                )
            }
            ("bam", Some(matches)) => {
                bam::collapse_reads_to_fragments::call_consensus_reads_from_paths(
                    matches.value_of("bam").unwrap(),
                    matches.value_of("consensus-bam").unwrap(),
                    value_t!(matches, "max-seq-dist", usize).unwrap(),
                    matches.is_present("verbose-read-names"),
                )
            }
            _ => unreachable!(),
        },
        ("sequence-stats", Some(matches)) => sequences_stats::stats(matches.is_present("fastq")),
        // This cannot be reached, since the matches step of
        // clap assures that a valid subcommand is provided
        _ => unreachable!(),
    }
}
<commit_msg>Create output path if it doesn't exist (#68)<commit_after>//! Documentation for Rust Bio Tools
use clap::{load_yaml, value_t};
use log::LevelFilter;

use clap::App;
use itertools::Itertools;
use std::collections::HashMap;
use std::error::Error;
use std::fs;
use std::path::Path;
use std::str::FromStr;

pub mod bam;
pub mod bcf;
pub mod common;
pub mod fastq;
pub mod sequences_stats;

fn main() -> Result<(), Box<dyn Error>> {
    let yaml = load_yaml!("cli.yaml");
    let matches = App::from_yaml(yaml)
        .version(env!("CARGO_PKG_VERSION"))
        .get_matches();

    fern::Dispatch::new()
        .format(|out, message, _| out.finish(format_args!("{}", message)))
        .level(if matches.is_present("verbose") {
            LevelFilter::Debug
        } else {
            LevelFilter::Info
        })
        .chain(std::io::stderr())
        .apply()
        .unwrap();

    match matches.subcommand() {
        ("fastq-split", Some(matches)) => {
            fastq::split::split(&matches.values_of("chunks").unwrap().collect_vec())
        }
        ("fastq-filter", Some(matches)) => fastq::filter::filter(&matches.value_of("ids").unwrap()),
        ("bam-depth", Some(matches)) => bam::depth::depth(
            &matches.value_of("bam-path").unwrap(),
            value_t!(matches, "max-read-length", u32).unwrap_or(1000),
            value_t!(matches, "include-flags", u16).unwrap_or(0),
            value_t!(matches, "exclude-flags", u16).unwrap_or(4 | 256 | 512 | 1024),
            value_t!(matches, "min-mapq", u8).unwrap_or(0),
        ),
        ("vcf-to-txt", Some(matches)) => bcf::to_txt::to_txt(
            &matches
                .values_of("info")
                .map(|values| values.collect_vec())
                .unwrap_or_default(),
            &matches
                .values_of("format")
                .map(|values| values.collect_vec())
                .unwrap_or_default(),
            matches.is_present("genotypes"),
        ),
        ("vcf-match", Some(matches)) => bcf::match_variants::match_variants(
            matches.value_of("vcf").unwrap(),
            value_t!(matches, "max-dist", u32).unwrap_or(20),
            value_t!(matches, "max-len-diff", u32).unwrap_or(10),
        ),
        ("vcf-baf", Some(_)) => bcf::baf::calculate_baf(),
        ("vcf-fix-iupac-alleles", Some(_)) => bcf::fix_iupac_alleles::fix_iupac_alleles(),
        ("vcf-annotate-dgidb", Some(matches)) => bcf::annotate_dgidb::annotate_dgidb(
            &matches.value_of("vcf").unwrap(),
            matches.value_of("api-path").unwrap().to_string(),
            &matches.value_of("field").unwrap(),
            match matches.is_present("datasources") {
                true => Some(matches.values_of("datasources").unwrap().collect()),
                false => None,
            },
            value_t!(matches, "genes-per-request", usize).unwrap(),
        ),
        ("vcf-report", Some(matches)) => {
            let mut sample_calls = HashMap::new();
            let mut bam_paths = HashMap::new();
            let output_path = matches.value_of("output-path").unwrap();
            if !Path::new(output_path).exists() {
                fs::create_dir(Path::new(output_path)).unwrap_or_else(|_| {
                    panic!(
                        "Couldn't create output directory at {}. Please make sure the path exists.",
                        output_path
                    )
                });
            }
            let max_cells = u32::from_str(matches.value_of("max-cells").unwrap()).unwrap();
            let custom_js = matches.value_of("custom-js");
            bcf::report::embed_js(output_path, custom_js)?;
            bcf::report::embed_css(output_path)?;
            bcf::report::embed_html(output_path)?;
            let fasta_path = matches.value_of("fasta").unwrap();
            let detail_path = output_path.to_owned() + "/details/";
            fs::create_dir(Path::new(&detail_path))?;
            let vcfs = matches.values_of("vcfs").unwrap();
            let bams = matches.values_of("bams").unwrap();
            let infos = matches.values_of("info");
            let formats = matches.values_of("format");
            for vcf in vcfs {
                let v: Vec<_> = vcf.split('=').collect();
                match sample_calls.insert(v[0].to_owned(), v[1].to_owned()) {
                    None => {}
                    _ => panic!("Found duplicate sample name {}. Please make sure the provided sample names are unique.", v[0].to_owned())
                }
            }
            for bam in bams {
                let b: Vec<_> = bam.split('=').collect();
                let c: Vec<_> = b[0].split(':').collect();
                let rec = bam_paths.entry(c[0].to_owned()).or_insert_with(Vec::new);
                rec.push((c[1].to_owned(), b[1].to_owned()))
            }

            for sample in sample_calls.keys().sorted() {
                bcf::report::table_report::table_report(
                    sample_calls.get(sample).unwrap(),
                    fasta_path,
                    bam_paths
                        .get(sample)
                        .unwrap_or_else(|| panic!("No bam provided for sample {}", sample)),
                    output_path,
                    sample,
                    infos.clone(),
                    formats.clone(),
                )?;
            }

            bcf::report::oncoprint::oncoprint(&sample_calls, output_path, max_cells)
        }
        ("collapse-reads-to-fragments", Some(matches)) => match matches.subcommand() {
            ("fastq", Some(matches)) => {
                fastq::collapse_reads_to_fragments::call_consensus_reads_from_paths(
                    matches.value_of("fq1").unwrap(),
                    matches.value_of("fq2").unwrap(),
                    matches.value_of("consensus-fq1").unwrap(),
                    matches.value_of("consensus-fq2").unwrap(),
                    matches.value_of("consensus-fq3"),
                    value_t!(matches, "umi-len", usize).unwrap(),
                    value_t!(matches, "max-seq-dist", usize).unwrap(),
                    value_t!(matches, "max-umi-dist", usize).unwrap(),
                    matches.is_present("umi-on-reverse"),
                    matches.is_present("verbose-read-names"),
                    if matches.is_present("insert-size") {
                        Some(value_t!(matches, "insert-size", usize).unwrap())
                    } else {
                        None
                    },
                    if matches.is_present("std-dev") {
                        Some(value_t!(matches, "std-dev", usize).unwrap())
                    } else {
                        None
                    },
                )
            }
            ("bam", Some(matches)) => {
                bam::collapse_reads_to_fragments::call_consensus_reads_from_paths(
                    matches.value_of("bam").unwrap(),
                    matches.value_of("consensus-bam").unwrap(),
                    value_t!(matches, "max-seq-dist", usize).unwrap(),
                    matches.is_present("verbose-read-names"),
                )
            }
            _ => unreachable!(),
        },
        ("sequence-stats", Some(matches)) => sequences_stats::stats(matches.is_present("fastq")),
        // This cannot be reached, since the matches step of
        // clap assures that a valid subcommand is provided
        _ => unreachable!(),
    }
}
<|endoftext|>


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

extern crate libc;

mod args;
mod ansi;
mod matching;
mod screen;

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

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

  if args.help { return; }

  let choices = read_choices();
  let mut search = args.initial_search.clone();
  if args.use_first {
    let matches = matching::compute_matches(&choices, search.as_slice());
    println!("{}", matches[0]);
    return;
  }

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

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

  let start_line = screen.height - visible_choices - 1;
  loop {
    let mut matches = matching::compute_matches(&choices, search.as_slice());
    screen.hide_cursor();
    screen.blank_screen(start_line);
    screen.move_cursor(start_line, 0);
    let mut i = 1;
    screen.write(format!("> {} ({} choices)\n", search.as_slice(), choices.len()).as_slice());
    for choice in matches.iter() {
      if i == index + 1 {
        screen.write_inverted(choice.as_slice());
      } else {
        screen.write(choice.as_slice());
      }
      if i >= visible_choices as usize {
        break;
      } else {
        screen.write("\n");
      }
      i += 1;
    }

    screen.move_cursor(start_line, 2 + search.len() as u16);
    screen.show_cursor();

    let chars = screen.tty.get_buffered_keys();
    let mut matches_stale = false;
    for char in chars.iter() {
      match *char {
        Char(x) => {
          search.push(x);
          index = 0;
          matches_stale = true;
        }
        Backspace => { search.pop(); matches_stale = true; }
        Control('h') => { search.pop(); matches_stale = true; }
        Control('u') => { search.clear(); matches_stale = true; }
        Control('c') => { return; }
        Control('n') => { index = min(index + 1, min(visible_choices as usize - 1, matches.len() - 1)); }
        Control('p') => { index = if index == 0 { 0 } else { index - 1 }; }
        Enter => {
          screen.move_cursor(start_line + visible_choices, 0);
          screen.write("\n");
          if matches_stale {
            matches = matching::compute_matches(&choices, search.as_slice());
          }
          println!("{}", matches[index]);
          return;
        }
        _ => panic!("Unexpected input"),
      }
    }
  }
}

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

  while let Ok(mut s) = stdin.read_line() {
    trim(&mut s);
    lines.push(s);
  }

  lines
}

fn trim(s: &mut String) {
  while let Some(x) = s.pop() {
    if x != '\n' {
      s.push(x);
      return;
    }
  }
}
<commit_msg>Don't recompute matches on C-n/C-p<commit_after>#![allow(unstable, dead_code)]

extern crate libc;

mod args;
mod ansi;
mod matching;
mod screen;

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

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

  if args.help { return; }

  let choices = read_choices();
  let mut search = args.initial_search.clone();
  if args.use_first {
    let matches = matching::compute_matches(&choices, search.as_slice());
    println!("{}", matches[0]);
    return;
  }

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

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

  let start_line = screen.height - visible_choices - 1;
  let mut matches_stale = true;
  let mut matches = matching::compute_matches(&choices, search.as_slice());
  loop {
    if matches_stale {
      matches = matching::compute_matches(&choices, search.as_slice());
      matches_stale = false;
    }
    screen.hide_cursor();
    screen.blank_screen(start_line);
    screen.move_cursor(start_line, 0);
    let mut i = 1;
    screen.write(format!("> {} ({} choices)\n", search.as_slice(), choices.len()).as_slice());
    for choice in matches.iter() {
      if i == index + 1 {
        screen.write_inverted(choice.as_slice());
      } else {
        screen.write(choice.as_slice());
      }
      if i >= visible_choices as usize {
        break;
      } else {
        screen.write("\n");
      }
      i += 1;
    }

    screen.move_cursor(start_line, 2 + search.len() as u16);
    screen.show_cursor();

    let chars = screen.tty.get_buffered_keys();
    for char in chars.iter() {
      match *char {
        Char(x) => {
          search.push(x);
          index = 0;
          matches_stale = true;
        }
        Backspace => { search.pop(); matches_stale = true; }
        Control('h') => { search.pop(); matches_stale = true; }
        Control('u') => { search.clear(); matches_stale = true; }
        Control('c') => { return; }
        Control('n') => { index = min(index + 1, min(visible_choices as usize - 1, matches.len() - 1)); }
        Control('p') => { index = if index == 0 { 0 } else { index - 1 }; }
        Enter => {
          screen.move_cursor(start_line + visible_choices, 0);
          screen.write("\n");
          if matches_stale {
            matches = matching::compute_matches(&choices, search.as_slice());
          }
          println!("{}", matches[index]);
          return;
        }
        _ => panic!("Unexpected input"),
      }
    }
  }
}

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

  while let Ok(mut s) = stdin.read_line() {
    trim(&mut s);
    lines.push(s);
  }

  lines
}

fn trim(s: &mut String) {
  while let Some(x) = s.pop() {
    if x != '\n' {
      s.push(x);
      return;
    }
  }
}
<|endoftext|>


<commit_before>// Worklog is for recording your hours.
//

extern crate chrono;
extern crate csv;
extern crate clap;
extern crate rustc_serialize;

mod error;
mod timeclock;
mod util;

use chrono::*;
use clap::{Arg, ArgGroup, App};
use error::WorklogError;
use std::env;
use std::fs::OpenOptions;
use std::io::prelude::*;
use std::path::PathBuf;
use timeclock::Direction;
use timeclock::now;

#[cfg(target_family = "unix")]
static CSV_FILE_NAME: &'static str = ".worklog.csv";
#[cfg(target_family = "windows")]
static CSV_FILE_NAME: &'static str = "worklog.csv";

// TODO This belongs in a config file somewhere
const WEEKSTART: i64 = Weekday::Sat as i64;

fn print_csv_entries<R: Read>(file: R) -> Result<(), WorklogError> {
    let csv_entries = try!(timeclock::read_timesheet(file));
    for rec in csv_entries {
        println!("{}", rec);
    }
    Ok(())
}

fn print_full_summary<R: Read>(file: R,
                               rounding: util::Rounding)
                               -> Result<(), WorklogError> {
    let csv_entries = try!(timeclock::read_timesheet(file));
    let records = timeclock::collect_date_records(csv_entries);

    let mut total_hours: f64 = 0.0;
    for rec in records {
        let hours = util::round(rec.seconds(), rounding) / 3600.0;
        total_hours += hours;
        println!("{} {:.2} {}", rec.date().format("%F"), hours, rec.memo());
    }

    println!("Total Hours: {:.2}", total_hours);
    Ok(())
}

fn print_short_summary<R: Read>(file: R,
                                since: Date<FixedOffset>,
                                rounding: util::Rounding)
                                -> Result<(), WorklogError> {
    let csv_entries = try!(timeclock::read_timesheet(file));
    let records = timeclock::collect_date_records(csv_entries);

    let mut total_hours: f64 = 0.0;
    for rec in records {
        if rec.date() >= since {
            let hours = util::round(rec.seconds(), rounding) / 3600.0;
            total_hours += hours;
            println!("{} {:.2} {}", rec.date().format("%F"), hours, rec.memo());
        }
    }

    println!("Total Hours: {:.2}", total_hours);
    Ok(())
}

/// Get the path for the csv data file
fn get_csv_path() -> Result<PathBuf, WorklogError> {
    let mut data_path = match env::home_dir() {
        Some(path) => path,
        None => {
            return Err(env::VarError::NotPresent).map_err(WorklogError::Env)
        }
    };

    data_path.push(CSV_FILE_NAME);
    Ok(data_path)
}

fn main0() -> Result<(), WorklogError> {
    // Using std env macro rather than depending on clap's. No difference
    // as far as I can tell.
    let matches = App::new(env!("CARGO_PKG_NAME"))
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .arg(Arg::from_usage("[in] -i, --in 'Record an In entry'"))
        .arg(Arg::from_usage("[out] -o, --out 'Record an Out entry'"))
        .arg(Arg::from_usage("[time] -t, --time <TIME> 'time'")
            .default_value("now")
            .requires("inout"))
        .arg(Arg::from_usage("[memo] -m, --memo <MEMO> 'Memo for the entry'")
            .default_value("")
            .requires("inout"))
        .arg(Arg::from_usage("[summary] -s, --summary 'Print a summary'")
            .conflicts_with("inout"))
        .arg(Arg::from_usage("[log] -l, --log 'Print the full log'")
            .conflicts_with("inout"))
        .arg(Arg::from_usage("[round] -r, --round 'Round totals up to the next quarter hour'")
            .conflicts_with("inout"))
        .group(ArgGroup::with_name("inout").args(&["in", "out"]))
        .get_matches();

    let csv_path = try!(get_csv_path());

    let mut csv_file = try!(OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .open(csv_path));

    let rounding = {
        if matches.is_present("round") {
            // round up in 15min intervals
            util::Rounding::Up(900.0)
        } else {
            util::Rounding::None
        }
    };

    if matches.is_present("in") || matches.is_present("out") {
        let dir = if matches.is_present("in") {
            Direction::In
        } else {
            Direction::Out
        };

        let time = matches.value_of("time").unwrap();
        let time = try!(util::parse_multi_time_fmt(&time));

        let memo = matches.value_of("memo").unwrap().to_owned();
        timeclock::mark_time(dir, time, memo, &mut csv_file);

        println!("Clocked {:#} at {}", dir, time.format("%F %I:%M %P"));

    } else if matches.is_present("summary") {
        try!(print_full_summary(&csv_file, rounding));

    } else if matches.is_present("log") {
        try!(print_csv_entries(&csv_file));

    } else {
        let days_back = (-WEEKSTART + 7) % 7;
        let start_date = now().date() - Duration::days(days_back);
        try!(print_short_summary(&csv_file, start_date, rounding));
    }

    Ok(())
}

fn main() {
    match main0() {
        Ok(_) => {}
        Err(err) => {
            let _ = writeln!(&mut std::io::stderr(), "Error: {}", err);
        }
    };
}
<commit_msg>Fixed incorrect date calculation<commit_after>// Worklog is for recording your hours.
//

extern crate chrono;
extern crate csv;
extern crate clap;
extern crate rustc_serialize;

mod error;
mod timeclock;
mod util;

use chrono::*;
use clap::{Arg, ArgGroup, App};
use error::WorklogError;
use std::env;
use std::fs::OpenOptions;
use std::io::prelude::*;
use std::path::PathBuf;
use timeclock::Direction;
use timeclock::now;

#[cfg(target_family = "unix")]
static CSV_FILE_NAME: &'static str = ".worklog.csv";
#[cfg(target_family = "windows")]
static CSV_FILE_NAME: &'static str = "worklog.csv";

// TODO This belongs in a config file somewhere
const WEEKSTART: i64 = Weekday::Sat as i64;

fn print_csv_entries<R: Read>(file: R) -> Result<(), WorklogError> {
    let csv_entries = try!(timeclock::read_timesheet(file));
    for rec in csv_entries {
        println!("{}", rec);
    }
    Ok(())
}

fn print_full_summary<R: Read>(file: R,
                               rounding: util::Rounding)
                               -> Result<(), WorklogError> {
    let csv_entries = try!(timeclock::read_timesheet(file));
    let records = timeclock::collect_date_records(csv_entries);

    let mut total_hours: f64 = 0.0;
    for rec in records {
        let hours = util::round(rec.seconds(), rounding) / 3600.0;
        total_hours += hours;
        println!("{} {:.2} {}", rec.date().format("%F"), hours, rec.memo());
    }

    println!("Total Hours: {:.2}", total_hours);
    Ok(())
}

fn print_short_summary<R: Read>(file: R,
                                since: Date<FixedOffset>,
                                rounding: util::Rounding)
                                -> Result<(), WorklogError> {
    let csv_entries = try!(timeclock::read_timesheet(file));
    let records = timeclock::collect_date_records(csv_entries);

    let mut total_hours: f64 = 0.0;
    for rec in records {
        if rec.date() >= since {
            let hours = util::round(rec.seconds(), rounding) / 3600.0;
            total_hours += hours;
            println!("{} {:.2} {}", rec.date().format("%F"), hours, rec.memo());
        }
    }

    println!("Total Hours: {:.2}", total_hours);
    Ok(())
}

/// Get the path for the csv data file
fn get_csv_path() -> Result<PathBuf, WorklogError> {
    let mut data_path = match env::home_dir() {
        Some(path) => path,
        None => {
            return Err(env::VarError::NotPresent).map_err(WorklogError::Env)
        }
    };

    data_path.push(CSV_FILE_NAME);
    Ok(data_path)
}

fn main0() -> Result<(), WorklogError> {
    // Using std env macro rather than depending on clap's. No difference
    // as far as I can tell.
    let matches = App::new(env!("CARGO_PKG_NAME"))
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .arg(Arg::from_usage("[in] -i, --in 'Record an In entry'"))
        .arg(Arg::from_usage("[out] -o, --out 'Record an Out entry'"))
        .arg(Arg::from_usage("[time] -t, --time <TIME> 'time'")
            .default_value("now")
            .requires("inout"))
        .arg(Arg::from_usage("[memo] -m, --memo <MEMO> 'Memo for the entry'")
            .default_value("")
            .requires("inout"))
        .arg(Arg::from_usage("[summary] -s, --summary 'Print a summary'")
            .conflicts_with("inout"))
        .arg(Arg::from_usage("[log] -l, --log 'Print the full log'")
            .conflicts_with("inout"))
        .arg(Arg::from_usage("[round] -r, --round 'Round totals up to the next quarter hour'")
            .conflicts_with("inout"))
        .group(ArgGroup::with_name("inout").args(&["in", "out"]))
        .get_matches();

    let csv_path = try!(get_csv_path());

    let mut csv_file = try!(OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .open(csv_path));

    let rounding = {
        if matches.is_present("round") {
            // round up in 15min intervals
            util::Rounding::Up(900.0)
        } else {
            util::Rounding::None
        }
    };

    if matches.is_present("in") || matches.is_present("out") {
        let dir = if matches.is_present("in") {
            Direction::In
        } else {
            Direction::Out
        };

        let time = matches.value_of("time").unwrap();
        let time = try!(util::parse_multi_time_fmt(&time));

        let memo = matches.value_of("memo").unwrap().to_owned();
        timeclock::mark_time(dir, time, memo, &mut csv_file);

        println!("Clocked {:#} at {}", dir, time.format("%F %I:%M %P"));

    } else if matches.is_present("summary") {
        try!(print_full_summary(&csv_file, rounding));

    } else if matches.is_present("log") {
        try!(print_csv_entries(&csv_file));

    } else {
        let today = now().date();
        let weekday = today.weekday() as i64;
        let days_back = (7 - WEEKSTART + weekday) % 7;
        let start_date = today - Duration::days(days_back);
        try!(print_short_summary(&csv_file, start_date, rounding));
    }

    Ok(())
}

fn main() {
    match main0() {
        Ok(_) => {}
        Err(err) => {
            let _ = writeln!(&mut std::io::stderr(), "Error: {}", err);
        }
    };
}
<|endoftext|>


<commit_before>// Genetic Algorithm for convergence to a given string
// Author: Duncan Dean
extern crate rand;
use rand::Rng;

fn rand_string(size: u32) -> String {
    let mut rand_string = "".to_string();
    for _ in 0..size {
        rand_string = rand_string + &(rand::thread_rng().gen_range(32, 123) as u8 as char).to_string();
    }
    rand_string
}

fn crossover(chrome1: Chromosome, chrome2: Chromosome, cross_prob: f32, solution: &str) -> Chromosome {
    if rand::thread_rng().gen_range::<f32>(0.0, 1.0) <= cross_prob {
        let code_len = chrome1.code.len();
        let rand_index = rand::thread_rng().gen_range(0, code_len);
        let first_half = &chrome1.code[0..rand_index];
        let second_half = &chrome2.code[rand_index..code_len];
        let new_code = first_half.to_string() + second_half;
        Chromosome {
            code: new_code.to_string(),
            cost_score: cost_function(&new_code, solution),
            solution: solution.to_string(),
        }
    } else if chrome1.cost_score < chrome2.cost_score {
        chrome1.clone()
    } else {
        chrome2.clone()
    }
}

fn cost_function(code: &str, solution: &str) -> u32 {
    let chars: Vec<char> = solution.to_string().chars().collect();
    let code_chars: Vec<char> = code.to_string().chars().collect();
    let mut cost: u32 = 0;
    for i in 0..chars.len() {
        cost += ((chars[i] as u8 as i32) - (code_chars[i] as u8 as i32)).abs() as u32;
    }
    cost
}

#[derive(Clone)]
struct Chromosome {
    code: String,
    cost_score: u32,
    solution: String,
}

impl Chromosome {
    fn new(solution: String) -> Chromosome {
        let random_string = &rand_string(solution.len() as u32);
        Chromosome {
            code: random_string.to_string(),
            cost_score: cost_function(random_string, &solution),
            solution: solution,
        }
    }

    fn mutate(&mut self, mut_prob:f64) -> Chromosome {
        if rand::thread_rng().gen_range::<f64>(0.0, 1.0) <= mut_prob {
            let rand_index = rand::thread_rng().gen_range(0, self.code.len());
            let mut new_code_vec: Vec<char> = self.code.chars().collect();
            new_code_vec[rand_index] = rand::thread_rng().gen_range(32, 123) as u8 as char;
            self.code = new_code_vec.iter().fold("".to_string(), |acc, s| acc + &s.to_string());
            self.cost_score = cost_function(&self.code, &self.solution);
            self.clone()

        } else {
            self.clone()
        }

    }

}

fn main() {
    let solution = "The quick brown fox jumps over the lazy dog.".to_string();
    let gen_number = 600;
    let mut pop_size = 40; // Must be greater than 1
    let max_pop = pop_size * 10;
    let mut_prob = 0.4;
    let kill_constant = 0.45;
    let cross_prob = 0.6;
    let mut population = Vec::new();

    println!("{}", solution.len());
    // Initialize population
    for i in 0..pop_size {
        population.push(Chromosome::new(solution.to_string()));
        println!("{}, cost: {}", population[i].code, population[i].cost_score);
    }

    // Step generation
    for _ in 1..(gen_number + 1) {
        // Mutations
        for i in 0..pop_size {
            population[i] = population[i].mutate(mut_prob);
            //println!("{}, cost: {}", population[i].code, population[i].cost_score);
        }

        // Kill the weaklings
        population.sort_by(|ref a, ref b| a.cost_score.cmp(&b.cost_score));
        for _ in 0..((( pop_size as f32)*kill_constant) as u32) {
            population.pop();
        }
        pop_size = population.len();

        // Crossovers
        if pop_size < max_pop {
            for i in 0..(population.len() - 2) {
                let chrome1 = population[i].clone();
                let chrome2 = population[i+1].clone();
                population.push(crossover(chrome1, chrome2, cross_prob, &solution));
                pop_size += 1;
            }
        }

    }
    println!("\n\nFinal");
    for i in 0..pop_size {
        println!("{}, cost: {}", population[i].code, population[i].cost_score);
    }

}
<commit_msg>Update<commit_after>// Genetic Algorithm for convergence to a given string
// Author: Duncan Dean
extern crate rand;
use rand::Rng;

fn rand_string(size: u32) -> String {
    let mut rand_string = String::with_capacity(size as usize);
    for _ in 0..size {
        rand_string = rand_string + &(rand::thread_rng().gen_range(32, 123) as u8 as char).to_string();
    }
    rand_string
}

fn crossover(chrome1: Chromosome, chrome2: Chromosome, cross_prob: f32, solution: &str) -> Chromosome {
    if rand::thread_rng().gen_range::<f32>(0.0, 1.0) <= cross_prob {
        let code_len = chrome1.code.len();
        let rand_index = rand::thread_rng().gen_range(0, code_len);
        let first_half = &chrome1.code[0..rand_index];
        let second_half = &chrome2.code[rand_index..code_len];
        let new_code = first_half.to_string() + second_half;
        Chromosome {
            code: new_code.to_string(),
            cost_score: cost_function(&new_code, solution),
            solution: solution.to_string(),
        }
    } else if chrome1.cost_score < chrome2.cost_score {
        chrome1.clone()
    } else {
        chrome2.clone()
    }
}

fn cost_function(code: &str, solution: &str) -> u32 {
    let chars: Vec<char> = solution.to_string().chars().collect();
    let code_chars: Vec<char> = code.to_string().chars().collect();
    let mut cost: u32 = 0;
    for i in 0..chars.len() {
        cost += ((chars[i] as u8 as i32) - (code_chars[i] as u8 as i32)).abs() as u32;
    }
    cost
}

#[derive(Clone)]
struct Chromosome {
    code: String,
    cost_score: u32,
    solution: String,
}

impl Chromosome {
    fn new(solution: String) -> Chromosome {
        let random_string = &rand_string(solution.len() as u32);
        Chromosome {
            code: random_string.to_string(),
            cost_score: cost_function(random_string, &solution),
            solution: solution,
        }
    }

    fn mutate(&mut self, mut_prob:f64) -> Chromosome {
        if rand::thread_rng().gen_range::<f64>(0.0, 1.0) <= mut_prob {
            let rand_index = rand::thread_rng().gen_range(0, self.code.len());
            let mut new_code_vec: Vec<char> = self.code.chars().collect();
            new_code_vec[rand_index] = rand::thread_rng().gen_range(32, 123) as u8 as char;
            self.code = new_code_vec.iter().fold("".to_string(), |acc, s| acc + &s.to_string());
            self.cost_score = cost_function(&self.code, &self.solution);
            self.clone()

        } else {
            self.clone()
        }

    }

}

fn main() {
    let solution = "The quick brown fox jumps over the lazy dog.".to_string();
    let gen_number = 600;
    let pop_size = 30; // Must be greater than 1
    let max_pop = pop_size * 5;
    let mut_prob = 0.4;
    let kill_constant = 0.45;
    let cross_prob = 0.6;
    let mut population = Vec::new();

    // Initialize population
    for i in 0..pop_size {
        population.push(Chromosome::new(solution.to_string()));
        println!("{}, cost: {}", population[i].code, population[i].cost_score);
    }

    population.sort_by(|ref a, ref b| a.cost_score.cmp(&b.cost_score));
    let mut winner: Chromosome = population[0].clone();
    // Step generation
    let mut j = 1;
    while  winner.cost_score != 0 {

        // Mutations
        for i in 0..population.len() {
            population[i] = population[i].mutate(mut_prob);
        }

        // Kill the weaklings
        population.sort_by(|ref a, ref b| a.cost_score.cmp(&b.cost_score));
        for _ in 0..((( population.len() as f32)*kill_constant) as u32) {
            population.pop();
        }

        // Crossovers
        if population.len() < max_pop {
            for i in 0..(population.len() - 2) {
                let chrome1 = population[i].clone();
                let chrome2 = population[i+1].clone();
                population.push(crossover(chrome1, chrome2, cross_prob, &solution));
            }
        }

        population.sort_by(|ref a, ref b| a.cost_score.cmp(&b.cost_score));
        winner = population[0].clone();
        println!("\nGeneration: {}\nString: {}\nCost: {}\n", j, winner.code, winner.cost_score);

        j += 1;

    }

}
<|endoftext|>


<commit_before>#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate clap;

extern crate qrcode;
extern crate termion;
extern crate image;
extern crate regex;

use clap::{App, Arg, SubCommand, AppSettings};
use qrcode::{QrCode, EcLevel};
use termion::color;
use image::GrayImage;

mod payloads;

const WIFI_COMMAND: &'static str = "wifi";
const MAIL_COMMAND: &'static str = "mail";
// const SMS_COMMAND: &'static str = "sms";
// const MMS_COMMAND: &'static str = "mms";
// const GEO_COMMAND: &'static str = "geo";
const PHONE_COMMAND: &'static str = "phone";
// const SKYPE_COMMAND: &'static str = "skype";
const URL_COMMAND: &'static str = "url";
// const BOOKMARK_COMMAND: &'static str = "bookmark";
// const BITCOIN_COMMAND: &'static str = "bitcoin";
// const GIRO_COMMAND: &'static str = "giro";
// const CALENDAR_COMMAND: &'static str = "calendar";

/*TODO: add arguments for:
- Add help texts for the subcommands
- Add error texts for some exit branches (mostly file output and qrcode gen)
- shell completions (see build.rasst file)
- implement different types of qr payloads
*/
fn main() {
    // create the cli interface with all subcommands, flags and args
    let app = App::new(crate_name!())
        .version(crate_version!())
        .author(crate_authors!("\n"))
        .about(crate_description!())
        .global_setting(AppSettings::SubcommandsNegateReqs)
        .arg(Arg::with_name("safe_zone")
            .global(true)
            .required(false)
            .short("s")
            .long("safe-zone")
            .help("Sets wether the safe zone around the code should be drawn or not.")
            .takes_value(false)
            .multiple(false))
        .arg(Arg::with_name("output")
            .global(true)
            .short("o").long("output")
            .help("Prints the QR-Code to a file. The image format is derived from the file extension. Currently only jpeg and png files are supported.")
            .takes_value(true)
            .value_name("FILE")
            .multiple(false))
        .arg(Arg::with_name("error")
            .global(true)
            .short("e").long("error")
            .help("Set the desired error correction level")
            .takes_value(true)
            .value_name("LEVEL")
            .possible_values(&["L", "M", "Q", "H"])
            .default_value("H"))
        .arg(Arg::with_name("INPUT")
            .help("The input string to use")
            .required(true))
        .subcommand(SubCommand::with_name(WIFI_COMMAND)
            .about("formats to a wifi access string QR-Code")
            .arg(Arg::with_name("ssid").required(true))
            .arg(Arg::with_name("pwd").required(true))
            .arg(Arg::with_name("mode")
                .possible_values(&["WEP", "WPA", "nopass"])
                .default_value("WPA")
                .value_name("MODE"))
            .arg(Arg::with_name("hidden")
                .required(false)
                .possible_values(&["true", "false"])
                .default_value("false")
                .value_name("HIDDEN")))
        .subcommand(SubCommand::with_name(MAIL_COMMAND)
            .about("formats to a mail adress string QR-Code")
            .arg(Arg::with_name("receiver").required(true))
            .arg(Arg::with_name("subject"))
            .arg(Arg::with_name("message"))
            .arg(Arg::with_name("encoding")
                .possible_values(&["MAILTO", "MATMSG", "SMTP"])
                .default_value("MAILTO")
                .value_name("ENCODING")))
        .subcommand(SubCommand::with_name(URL_COMMAND)
            .about("formats to an URL QR-Code")
            .arg(Arg::with_name("url")
                .required(true)
                .value_name("URL")))
        .subcommand(SubCommand::with_name(PHONE_COMMAND)
            .about("formats to a phone number QR-Code")
            .arg(Arg::with_name("number")
                .required(true)
                .value_name("NUMBER"))
        );

    let matches = app.get_matches();

    // deduce the string payload
    let payload = get_payload(&matches);

    // should we draw a white border (safe zone) around the code?
    let safe: bool = match matches.occurrences_of("safe_zone") {
        0 => true,
        _ => false,
    };

    // what error level can we expect? always defaults to "H"
    let error: EcLevel = match matches.value_of("error").unwrap() {
        "L" => EcLevel::L,
        "M" => EcLevel::M,
        "Q" => EcLevel::Q,
        &_ => EcLevel::H,
    };

    //TODO: catch the possible errors
    let code = QrCode::with_error_correction_level(payload, error).unwrap();

    // are we drawing to the terminal or to a file?
    match matches.occurrences_of("output") {
        1 => save(&code, safe, matches.value_of("output").unwrap()),
        _ => draw(&code, safe),
    }
}

// deduces wich kind of string we are going to encode
fn get_payload(matches: &clap::ArgMatches) -> String {
    if let Some(sub) = matches.subcommand_matches(WIFI_COMMAND) {
        let auth = match sub.value_of("mode") {
            Some("WEP") => payloads::Authentication::WEP,
            Some("WPA") => payloads::Authentication::WPA,
            _ => payloads::Authentication::nopass,
        };
        return payloads::wifi_string(sub.value_of("ssid").unwrap(),
                                     sub.value_of("pwd").unwrap(),
                                     &auth,
                                     sub.value_of("hidden").unwrap() == "true");
    } else if let Some(sub) = matches.subcommand_matches(MAIL_COMMAND) {
        let encoding = match sub.value_of("encoding") {
            Some("MATMSG") => payloads::MailEncoding::MATMSG,
            Some("SMTP") => payloads::MailEncoding::SMTP,
            _ => payloads::MailEncoding::MAILTO,
        };
        return payloads::mail_string(sub.value_of("receiver").unwrap(),
                                     sub.value_of("subject").unwrap(),
                                     sub.value_of("message").unwrap(),
                                     &encoding);
    } else if let Some(sub) = matches.subcommand_matches(URL_COMMAND) {
        return payloads::url_string(sub.value_of("url").unwrap());
    } else if let Some(sub) = matches.subcommand_matches(PHONE_COMMAND) {
        return payloads::phone_string(sub.value_of("phone").unwrap());
    } else {
        return String::from(matches.value_of("INPUT").unwrap());
    }
}

// save to a file at the path
fn save(code: &QrCode, safe: bool, path: &str) {
    // render to a image struct
    let image: GrayImage = code.render().quiet_zone(safe).to_image();

    // save the image
    match image.save(path) {
        Ok(..) => println!("Image successfully saved to: {:?}", path),
        Err(e) => println!("Tried to create file but there was a problem: {:?}",
            if let Some(inner_err) = e.get_ref() { inner_err.to_string() } else { e.to_string() }),
    };
}

// draw to the terminal
fn draw(code: &QrCode, safe: bool) {
    // get "bit" array
    let bit_array = code.to_vec();

    // get the code width and add extra space for the safe zone
    let w = code.width();
    let wide = w + 6;

    // draw the first white safe zone
    if safe {
        for a in 1..(wide * 3) + 1 {
            print!("{}  ", color::Bg(color::LightWhite));
            if a % wide == 0 {
                println!("{}", color::Bg(color::Reset));
            }
        }
    }

    // main drawing loop
    for (i, item) in bit_array.iter().enumerate() {
        // left safe zone
        if safe && i % w == 0 {
            print!("{}      ", color::Bg(color::LightWhite));
        }

        // draw black or white blocks
        if *item {
            print!("{}  ", color::Bg(color::Black));
        } else {
            print!("{}  ", color::Bg(color::LightWhite));
        }

        if (i + 1) % w == 0 {
            if safe {
                //draw right safe zone
                println!("{}      {}",
                         color::Bg(color::LightWhite),
                         color::Bg(color::Reset));
            } else {
                println!("{}", color::Bg(color::Reset));
            }
        }
    }

    //Draw the last white safe zone
    if safe {
        for a in 1..(wide * 3) + 1 {
            print!("{}  ", color::Bg(color::LightWhite));
            if a % wide == 0 {
                println!("{}", color::Bg(color::Reset));
            }
        }
    }
}
<commit_msg>removed unused argument paramter<commit_after>#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate clap;

extern crate qrcode;
extern crate termion;
extern crate image;
extern crate regex;

use clap::{App, Arg, SubCommand, AppSettings};
use qrcode::{QrCode, EcLevel};
use termion::color;
use image::GrayImage;

mod payloads;

const WIFI_COMMAND: &'static str = "wifi";
const MAIL_COMMAND: &'static str = "mail";
// const SMS_COMMAND: &'static str = "sms";
// const MMS_COMMAND: &'static str = "mms";
// const GEO_COMMAND: &'static str = "geo";
const PHONE_COMMAND: &'static str = "phone";
// const SKYPE_COMMAND: &'static str = "skype";
const URL_COMMAND: &'static str = "url";
// const BOOKMARK_COMMAND: &'static str = "bookmark";
// const BITCOIN_COMMAND: &'static str = "bitcoin";
// const GIRO_COMMAND: &'static str = "giro";
// const CALENDAR_COMMAND: &'static str = "calendar";

/*TODO: add arguments for:
- Add help texts for the subcommands
- Add error texts for some exit branches (mostly file output and qrcode gen)
- shell completions (see build.rasst file)
- implement different types of qr payloads
*/
fn main() {
    // create the cli interface with all subcommands, flags and args
    let app = App::new(crate_name!())
        .version(crate_version!())
        .author(crate_authors!("\n"))
        .about(crate_description!())
        .global_setting(AppSettings::SubcommandsNegateReqs)
        .arg(Arg::with_name("safe_zone")
            .global(true)
            .required(false)
            .short("s")
            .long("safe-zone")
            .help("Sets wether the safe zone around the code should be drawn or not.")
            .takes_value(false)
            .multiple(false))
        .arg(Arg::with_name("output")
            .global(true)
            .short("o").long("output")
            .help("Prints the QR-Code to a file. The image format is derived from the file extension. Currently only jpeg and png files are supported.")
            .takes_value(true)
            .value_name("FILE")
            .multiple(false))
        .arg(Arg::with_name("error")
            .global(true)
            .short("e").long("error")
            .help("Set the desired error correction level")
            .takes_value(true)
            .value_name("LEVEL")
            .possible_values(&["L", "M", "Q", "H"])
            .default_value("H"))
        .arg(Arg::with_name("INPUT")
            .help("The input string to use")
            .required(true))
        .subcommand(SubCommand::with_name(WIFI_COMMAND)
            .about("formats to a wifi access string QR-Code")
            .arg(Arg::with_name("ssid").required(true))
            .arg(Arg::with_name("pwd").required(true))
            .arg(Arg::with_name("mode")
                .possible_values(&["WEP", "WPA", "nopass"])
                .default_value("WPA")
                .value_name("MODE"))
            .arg(Arg::with_name("hidden")
                .possible_values(&["true", "false"])
                .default_value("false")
                .value_name("HIDDEN")))
        .subcommand(SubCommand::with_name(MAIL_COMMAND)
            .about("formats to a mail adress string QR-Code")
            .arg(Arg::with_name("receiver").required(true))
            .arg(Arg::with_name("subject"))
            .arg(Arg::with_name("message"))
            .arg(Arg::with_name("encoding")
                .possible_values(&["MAILTO", "MATMSG", "SMTP"])
                .default_value("MAILTO")
                .value_name("ENCODING")))
        .subcommand(SubCommand::with_name(URL_COMMAND)
            .about("formats to an URL QR-Code")
            .arg(Arg::with_name("url")
                .required(true)
                .value_name("URL")))
        .subcommand(SubCommand::with_name(PHONE_COMMAND)
            .about("formats to a phone number QR-Code")
            .arg(Arg::with_name("number")
                .required(true)
                .value_name("NUMBER"))
        );

    let matches = app.get_matches();

    // deduce the string payload
    let payload = get_payload(&matches);

    // should we draw a white border (safe zone) around the code?
    let safe: bool = match matches.occurrences_of("safe_zone") {
        0 => true,
        _ => false,
    };

    // what error level can we expect? always defaults to "H"
    let error: EcLevel = match matches.value_of("error").unwrap() {
        "L" => EcLevel::L,
        "M" => EcLevel::M,
        "Q" => EcLevel::Q,
        &_ => EcLevel::H,
    };

    //TODO: catch the possible errors
    let code = QrCode::with_error_correction_level(payload, error).unwrap();

    // are we drawing to the terminal or to a file?
    match matches.occurrences_of("output") {
        1 => save(&code, safe, matches.value_of("output").unwrap()),
        _ => draw(&code, safe),
    }
}

// deduces wich kind of string we are going to encode
fn get_payload(matches: &clap::ArgMatches) -> String {
    if let Some(sub) = matches.subcommand_matches(WIFI_COMMAND) {
        let auth = match sub.value_of("mode") {
            Some("WEP") => payloads::Authentication::WEP,
            Some("WPA") => payloads::Authentication::WPA,
            _ => payloads::Authentication::nopass,
        };
        return payloads::wifi_string(sub.value_of("ssid").unwrap(),
                                     sub.value_of("pwd").unwrap(),
                                     &auth,
                                     sub.value_of("hidden").unwrap() == "true");
    } else if let Some(sub) = matches.subcommand_matches(MAIL_COMMAND) {
        let encoding = match sub.value_of("encoding") {
            Some("MATMSG") => payloads::MailEncoding::MATMSG,
            Some("SMTP") => payloads::MailEncoding::SMTP,
            _ => payloads::MailEncoding::MAILTO,
        };
        return payloads::mail_string(sub.value_of("receiver").unwrap(),
                                     sub.value_of("subject").unwrap(),
                                     sub.value_of("message").unwrap(),
                                     &encoding);
    } else if let Some(sub) = matches.subcommand_matches(URL_COMMAND) {
        return payloads::url_string(sub.value_of("url").unwrap());
    } else if let Some(sub) = matches.subcommand_matches(PHONE_COMMAND) {
        return payloads::phone_string(sub.value_of("phone").unwrap());
    } else {
        return String::from(matches.value_of("INPUT").unwrap());
    }
}

// save to a file at the path
fn save(code: &QrCode, safe: bool, path: &str) {
    // render to a image struct
    let image: GrayImage = code.render().quiet_zone(safe).to_image();

    // save the image
    match image.save(path) {
        Ok(..) => println!("Image successfully saved to: {:?}", path),
        Err(e) => println!("Tried to create file but there was a problem: {:?}",
            if let Some(inner_err) = e.get_ref() { inner_err.to_string() } else { e.to_string() }),
    };
}

// draw to the terminal
fn draw(code: &QrCode, safe: bool) {
    // get "bit" array
    let bit_array = code.to_vec();

    // get the code width and add extra space for the safe zone
    let w = code.width();
    let wide = w + 6;

    // draw the first white safe zone
    if safe {
        for a in 1..(wide * 3) + 1 {
            print!("{}  ", color::Bg(color::LightWhite));
            if a % wide == 0 {
                println!("{}", color::Bg(color::Reset));
            }
        }
    }

    // main drawing loop
    for (i, item) in bit_array.iter().enumerate() {
        // left safe zone
        if safe && i % w == 0 {
            print!("{}      ", color::Bg(color::LightWhite));
        }

        // draw black or white blocks
        if *item {
            print!("{}  ", color::Bg(color::Black));
        } else {
            print!("{}  ", color::Bg(color::LightWhite));
        }

        if (i + 1) % w == 0 {
            if safe {
                //draw right safe zone
                println!("{}      {}",
                         color::Bg(color::LightWhite),
                         color::Bg(color::Reset));
            } else {
                println!("{}", color::Bg(color::Reset));
            }
        }
    }

    //Draw the last white safe zone
    if safe {
        for a in 1..(wide * 3) + 1 {
            print!("{}  ", color::Bg(color::LightWhite));
            if a % wide == 0 {
                println!("{}", color::Bg(color::Reset));
            }
        }
    }
}
<|endoftext|>