<commit_before>use std::env;
use std::path::PathBuf;

use cratesfyi::db::{self, add_path_into_database, connect_db};
use cratesfyi::utils::{add_crate_to_queue, remove_crate_priority, set_crate_priority};
use cratesfyi::Server;
use cratesfyi::{DocBuilder, DocBuilderOptions, RustwideBuilder};
use structopt::StructOpt;

pub fn main() {
    let _ = dotenv::dotenv();
    logger_init();

    CommandLine::from_args().handle_args();
}

fn logger_init() {
    use std::io::Write;

    let mut builder = env_logger::Builder::new();
    builder.format(|buf, record| {
        writeln!(
            buf,
            "{} [{}] {}: {}",
            time::now().strftime("%Y/%m/%d %H:%M:%S").unwrap(),
            record.level(),
            record.target(),
            record.args()
        )
    });
    builder.parse_filters(
        env::var("RUST_LOG")
            .ok()
            .as_deref()
            .unwrap_or("cratesfyi=info"),
    );

    rustwide::logging::init_with(builder.build());
}

#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
#[structopt(
    name = "cratesfyi",
    about = env!("CARGO_PKG_DESCRIPTION"),
    version = cratesfyi::BUILD_VERSION,
    rename_all = "kebab-case",
)]
enum CommandLine {
    Build(Build),

    /// Starts web server
    StartWebServer {
        #[structopt(name = "SOCKET_ADDR", default_value = "0.0.0.0:3000")]
        socket_addr: String,
    },

    /// Starts cratesfyi daemon
    Daemon {
        /// Run the server in the foreground instead of detaching a child
        #[structopt(name = "FOREGROUND", short = "f", long = "foreground")]
        foreground: bool,
    },

    /// Database operations
    Database {
        #[structopt(subcommand)]
        subcommand: DatabaseSubcommand,
    },

    /// Interactions with the build queue
    Queue {
        #[structopt(subcommand)]
        subcommand: QueueSubcommand,
    },
}

impl CommandLine {
    pub fn handle_args(self) {
        match self {
            Self::Build(build) => build.handle_args(),
            Self::StartWebServer { socket_addr } => {
                Server::start(Some(&socket_addr));
            }
            Self::Daemon { foreground } => cratesfyi::utils::start_daemon(!foreground),
            Self::Database { subcommand } => subcommand.handle_args(),
            Self::Queue { subcommand } => subcommand.handle_args(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
enum QueueSubcommand {
    /// Add a crate to the build queue
    Add {
        /// Name of crate to build
        #[structopt(name = "CRATE_NAME")]
        crate_name: String,
        /// Version of crate to build
        #[structopt(name = "CRATE_VERSION")]
        crate_version: String,
        /// Priority of build (new crate builds get priority 0)
        #[structopt(
            name = "BUILD_PRIORITY",
            short = "p",
            long = "priority",
            default_value = "5"
        )]
        build_priority: i32,
    },

    /// Interactions with build queue priorities
    DefaultPriority {
        #[structopt(subcommand)]
        subcommand: PrioritySubcommand,
    },
}

impl QueueSubcommand {
    pub fn handle_args(self) {
        match self {
            Self::Add {
                crate_name,
                crate_version,
                build_priority,
            } => {
                let conn = connect_db().expect("Could not connect to database");
                add_crate_to_queue(&conn, &crate_name, &crate_version, build_priority)
                    .expect("Could not add crate to queue");
            }

            Self::DefaultPriority { subcommand } => subcommand.handle_args(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
enum PrioritySubcommand {
    /// Set all crates matching a pattern to a priority level
    Set {
        /// See https://www.postgresql.org/docs/current/functions-matching.html for pattern syntax
        #[structopt(name = "PATTERN")]
        pattern: String,
        /// The priority to give crates matching the given `PATTERN`
        priority: i32,
    },

    /// Remove the prioritization of crates for a pattern
    Remove {
        /// See https://www.postgresql.org/docs/current/functions-matching.html for pattern syntax
        #[structopt(name = "PATTERN")]
        pattern: String,
    },
}

impl PrioritySubcommand {
    pub fn handle_args(self) {
        match self {
            Self::Set { pattern, priority } => {
                let conn = connect_db().expect("Could not connect to the database");

                set_crate_priority(&conn, &pattern, priority)
                    .expect("Could not set pattern's priority");
            }

            Self::Remove { pattern } => {
                let conn = connect_db().expect("Could not connect to the database");

                if let Some(priority) = remove_crate_priority(&conn, &pattern)
                    .expect("Could not remove pattern's priority")
                {
                    println!("Removed pattern with priority {}", priority);
                } else {
                    println!("Pattern did not exist and so was not removed");
                }
            }
        }
    }
}

/// Builds documentation in a chroot environment
#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
#[structopt(rename_all = "kebab-case")]
struct Build {
    #[structopt(
        name = "PREFIX",
        short = "P",
        long = "prefix",
        env = "CRATESFYI_PREFIX"
    )]
    prefix: PathBuf,

    /// Sets crates.io-index path
    #[structopt(name = "CRATES_IO_INDEX_PATH", long = "crates-io-index-path")]
    crates_io_index_path: Option<PathBuf>,

    /// Skips building documentation if documentation exists
    #[structopt(name = "SKIP_IF_EXISTS", short = "s", long = "skip")]
    skip_if_exists: bool,

    /// Skips building documentation if build log exists
    #[structopt(name = "SKIP_IF_LOG_EXISTS", long = "skip-if-log-exists")]
    skip_if_log_exists: bool,

    /// Keeps build directory after build.
    #[structopt(
        name = "KEEP_BUILD_DIRECTORY",
        short = "k",
        long = "keep-build-directory"
    )]
    keep_build_directory: bool,

    #[structopt(subcommand)]
    subcommand: BuildSubcommand,
}

impl Build {
    pub fn handle_args(self) {
        let docbuilder = {
            let mut doc_options = DocBuilderOptions::from_prefix(self.prefix);

            if let Some(crates_io_index_path) = self.crates_io_index_path {
                doc_options.crates_io_index_path = crates_io_index_path;
            }

            doc_options.skip_if_exists = self.skip_if_exists;
            doc_options.skip_if_log_exists = self.skip_if_log_exists;
            doc_options.keep_build_directory = self.keep_build_directory;

            doc_options
                .check_paths()
                .expect("The given paths were invalid");

            DocBuilder::new(doc_options)
        };

        self.subcommand.handle_args(docbuilder);
    }
}

#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
enum BuildSubcommand {
    /// Builds documentation of every crate
    World,

    /// Builds documentation for a crate
    Crate {
        /// Crate name
        #[structopt(
            name = "CRATE_NAME",
            required_unless("LOCAL"),
            requires("CRATE_VERSION")
        )]
        crate_name: String,

        /// Version of crate
        #[structopt(name = "CRATE_VERSION")]
        crate_version: String,

        /// Build a crate at a specific path
        #[structopt(short = "l", long = "local", conflicts_with_all(&["CRATE_NAME", "CRATE_VERSION"]))]
        local: Option<PathBuf>,
    },

    /// update the currently installed rustup toolchain
    UpdateToolchain {
        /// Update the toolchain only if no toolchain is currently installed
        #[structopt(name = "ONLY_FIRST_TIME", long = "only-first-time")]
        only_first_time: bool,
    },

    /// Adds essential files for the installed version of rustc
    AddEssentialFiles,

    /// Locks cratesfyi daemon to stop building new crates
    Lock,

    /// Unlocks cratesfyi daemon to continue building new crates
    Unlock,

    PrintOptions,
}

impl BuildSubcommand {
    pub fn handle_args(self, mut docbuilder: DocBuilder) {
        match self {
            Self::World => {
                docbuilder.load_cache().expect("Failed to load cache");

                let mut builder = RustwideBuilder::init().unwrap();
                builder
                    .build_world(&mut docbuilder)
                    .expect("Failed to build world");

                docbuilder.save_cache().expect("Failed to save cache");
            }

            Self::Crate {
                crate_name,
                crate_version,
                local,
            } => {
                docbuilder.load_cache().expect("Failed to load cache");
                let mut builder = RustwideBuilder::init().expect("failed to initialize rustwide");

                if let Some(path) = local {
                    builder
                        .build_local_package(&mut docbuilder, &path)
                        .expect("Building documentation failed");
                } else {
                    builder
                        .build_package(&mut docbuilder, &crate_name, &crate_version, None)
                        .expect("Building documentation failed");
                }

                docbuilder.save_cache().expect("Failed to save cache");
            }

            Self::UpdateToolchain { only_first_time } => {
                if only_first_time {
                    let conn = db::connect_db().unwrap();
                    let res = conn
                        .query("SELECT * FROM config WHERE name = 'rustc_version';", &[])
                        .unwrap();

                    if !res.is_empty() {
                        println!("update-toolchain was already called in the past, exiting");
                        return;
                    }
                }

                let mut builder = RustwideBuilder::init().unwrap();
                builder
                    .update_toolchain()
                    .expect("failed to update toolchain");
            }

            Self::AddEssentialFiles => {
                let mut builder = RustwideBuilder::init().unwrap();
                builder
                    .add_essential_files()
                    .expect("failed to add essential files");
            }

            Self::Lock => docbuilder.lock().expect("Failed to lock"),
            Self::Unlock => docbuilder.unlock().expect("Failed to unlock"),
            Self::PrintOptions => println!("{:?}", docbuilder.options()),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
enum DatabaseSubcommand {
    /// Run database migrations
    Migrate {
        /// The database version to migrate to
        #[structopt(name = "VERSION")]
        version: Option<i64>,
    },

    /// Updates github stats for crates.
    UpdateGithubFields,

    AddDirectory {
        /// Path of file or directory
        #[structopt(name = "DIRECTORY")]
        directory: PathBuf,
        /// Prefix of files in database
        #[structopt(name = "PREFIX", env = "CRATESFYI_PREFIX")]
        prefix: String,
    },

    /// Updates monthly release activity chart
    UpdateReleaseActivity,

    /// Updates search index
    UpdateSearchIndex,

    /// Removes a whole crate from the database
    DeleteCrate {
        /// Name of the crate to delete
        #[structopt(name = "CRATE_NAME")]
        crate_name: String,
    },

    /// Blacklist operations
    Blacklist {
        #[structopt(subcommand)]
        command: BlacklistSubcommand,
    },
}

impl DatabaseSubcommand {
    pub fn handle_args(self) {
        match self {
            Self::Migrate { version } => {
                let conn = connect_db().expect("failed to connect to the database");
                db::migrate(version, &conn).expect("Failed to run database migrations");
            }

            Self::UpdateGithubFields => {
                cratesfyi::utils::github_updater().expect("Failed to update github fields");
            }

            Self::AddDirectory { directory, prefix } => {
                let conn = db::connect_db().expect("failed to connect to the database");
                add_path_into_database(&conn, &prefix, directory)
                    .expect("Failed to add directory into database");
            }

            // FIXME: This is actually util command not database
            Self::UpdateReleaseActivity => cratesfyi::utils::update_release_activity()
                .expect("Failed to update release activity"),

            Self::UpdateSearchIndex => {
                let conn = db::connect_db().expect("failed to connect to the database");
                db::update_search_index(&conn).expect("Failed to update search index");
            }

            Self::DeleteCrate { crate_name } => {
                let conn = db::connect_db().expect("failed to connect to the database");
                db::delete_crate(&conn, &crate_name).expect("failed to delete the crate");
            }

            Self::Blacklist { command } => command.handle_args(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
enum BlacklistSubcommand {
    /// List all crates on the blacklist
    List,

    /// Add a crate to the blacklist
    Add {
        /// Crate name
        #[structopt(name = "CRATE_NAME")]
        crate_name: String,
    },

    /// Remove a crate from the blacklist
    Remove {
        /// Crate name
        #[structopt(name = "CRATE_NAME")]
        crate_name: String,
    },
}

impl BlacklistSubcommand {
    fn handle_args(self) {
        let conn = db::connect_db().expect("failed to connect to the database");

        match self {
            Self::List => {
                let crates =
                    db::blacklist::list_crates(&conn).expect("failed to list crates on blacklist");

                println!("{}", crates.join("\n"));
            }

            Self::Add { crate_name } => db::blacklist::add_crate(&conn, &crate_name)
                .expect("failed to add crate to blacklist"),

            Self::Remove { crate_name } => db::blacklist::remove_crate(&conn, &crate_name)
                .expect("failed to remove crate from blacklist"),
        }
    }
}
<commit_msg>Allow passing --local to build crates<commit_after>use std::env;
use std::path::PathBuf;

use cratesfyi::db::{self, add_path_into_database, connect_db};
use cratesfyi::utils::{add_crate_to_queue, remove_crate_priority, set_crate_priority};
use cratesfyi::Server;
use cratesfyi::{DocBuilder, DocBuilderOptions, RustwideBuilder};
use structopt::StructOpt;

pub fn main() {
    let _ = dotenv::dotenv();
    logger_init();

    CommandLine::from_args().handle_args();
}

fn logger_init() {
    use std::io::Write;

    let mut builder = env_logger::Builder::new();
    builder.format(|buf, record| {
        writeln!(
            buf,
            "{} [{}] {}: {}",
            time::now().strftime("%Y/%m/%d %H:%M:%S").unwrap(),
            record.level(),
            record.target(),
            record.args()
        )
    });
    builder.parse_filters(
        env::var("RUST_LOG")
            .ok()
            .as_deref()
            .unwrap_or("cratesfyi=info"),
    );

    rustwide::logging::init_with(builder.build());
}

#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
#[structopt(
    name = "cratesfyi",
    about = env!("CARGO_PKG_DESCRIPTION"),
    version = cratesfyi::BUILD_VERSION,
    rename_all = "kebab-case",
)]
enum CommandLine {
    Build(Build),

    /// Starts web server
    StartWebServer {
        #[structopt(name = "SOCKET_ADDR", default_value = "0.0.0.0:3000")]
        socket_addr: String,
    },

    /// Starts cratesfyi daemon
    Daemon {
        /// Run the server in the foreground instead of detaching a child
        #[structopt(name = "FOREGROUND", short = "f", long = "foreground")]
        foreground: bool,
    },

    /// Database operations
    Database {
        #[structopt(subcommand)]
        subcommand: DatabaseSubcommand,
    },

    /// Interactions with the build queue
    Queue {
        #[structopt(subcommand)]
        subcommand: QueueSubcommand,
    },
}

impl CommandLine {
    pub fn handle_args(self) {
        match self {
            Self::Build(build) => build.handle_args(),
            Self::StartWebServer { socket_addr } => {
                Server::start(Some(&socket_addr));
            }
            Self::Daemon { foreground } => cratesfyi::utils::start_daemon(!foreground),
            Self::Database { subcommand } => subcommand.handle_args(),
            Self::Queue { subcommand } => subcommand.handle_args(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
enum QueueSubcommand {
    /// Add a crate to the build queue
    Add {
        /// Name of crate to build
        #[structopt(name = "CRATE_NAME")]
        crate_name: String,
        /// Version of crate to build
        #[structopt(name = "CRATE_VERSION")]
        crate_version: String,
        /// Priority of build (new crate builds get priority 0)
        #[structopt(
            name = "BUILD_PRIORITY",
            short = "p",
            long = "priority",
            default_value = "5"
        )]
        build_priority: i32,
    },

    /// Interactions with build queue priorities
    DefaultPriority {
        #[structopt(subcommand)]
        subcommand: PrioritySubcommand,
    },
}

impl QueueSubcommand {
    pub fn handle_args(self) {
        match self {
            Self::Add {
                crate_name,
                crate_version,
                build_priority,
            } => {
                let conn = connect_db().expect("Could not connect to database");
                add_crate_to_queue(&conn, &crate_name, &crate_version, build_priority)
                    .expect("Could not add crate to queue");
            }

            Self::DefaultPriority { subcommand } => subcommand.handle_args(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
enum PrioritySubcommand {
    /// Set all crates matching a pattern to a priority level
    Set {
        /// See https://www.postgresql.org/docs/current/functions-matching.html for pattern syntax
        #[structopt(name = "PATTERN")]
        pattern: String,
        /// The priority to give crates matching the given `PATTERN`
        priority: i32,
    },

    /// Remove the prioritization of crates for a pattern
    Remove {
        /// See https://www.postgresql.org/docs/current/functions-matching.html for pattern syntax
        #[structopt(name = "PATTERN")]
        pattern: String,
    },
}

impl PrioritySubcommand {
    pub fn handle_args(self) {
        match self {
            Self::Set { pattern, priority } => {
                let conn = connect_db().expect("Could not connect to the database");

                set_crate_priority(&conn, &pattern, priority)
                    .expect("Could not set pattern's priority");
            }

            Self::Remove { pattern } => {
                let conn = connect_db().expect("Could not connect to the database");

                if let Some(priority) = remove_crate_priority(&conn, &pattern)
                    .expect("Could not remove pattern's priority")
                {
                    println!("Removed pattern with priority {}", priority);
                } else {
                    println!("Pattern did not exist and so was not removed");
                }
            }
        }
    }
}

/// Builds documentation in a chroot environment
#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
#[structopt(rename_all = "kebab-case")]
struct Build {
    #[structopt(
        name = "PREFIX",
        short = "P",
        long = "prefix",
        env = "CRATESFYI_PREFIX"
    )]
    prefix: PathBuf,

    /// Sets crates.io-index path
    #[structopt(name = "CRATES_IO_INDEX_PATH", long = "crates-io-index-path")]
    crates_io_index_path: Option<PathBuf>,

    /// Skips building documentation if documentation exists
    #[structopt(name = "SKIP_IF_EXISTS", short = "s", long = "skip")]
    skip_if_exists: bool,

    /// Skips building documentation if build log exists
    #[structopt(name = "SKIP_IF_LOG_EXISTS", long = "skip-if-log-exists")]
    skip_if_log_exists: bool,

    /// Keeps build directory after build.
    #[structopt(
        name = "KEEP_BUILD_DIRECTORY",
        short = "k",
        long = "keep-build-directory"
    )]
    keep_build_directory: bool,

    #[structopt(subcommand)]
    subcommand: BuildSubcommand,
}

impl Build {
    pub fn handle_args(self) {
        let docbuilder = {
            let mut doc_options = DocBuilderOptions::from_prefix(self.prefix);

            if let Some(crates_io_index_path) = self.crates_io_index_path {
                doc_options.crates_io_index_path = crates_io_index_path;
            }

            doc_options.skip_if_exists = self.skip_if_exists;
            doc_options.skip_if_log_exists = self.skip_if_log_exists;
            doc_options.keep_build_directory = self.keep_build_directory;

            doc_options
                .check_paths()
                .expect("The given paths were invalid");

            DocBuilder::new(doc_options)
        };

        self.subcommand.handle_args(docbuilder);
    }
}

#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
enum BuildSubcommand {
    /// Builds documentation of every crate
    World,

    /// Builds documentation for a crate
    Crate {
        /// Crate name
        #[structopt(
            name = "CRATE_NAME",
            required_unless("local"),
            requires("CRATE_VERSION")
        )]
        crate_name: Option<String>,

        /// Version of crate
        #[structopt(name = "CRATE_VERSION")]
        crate_version: Option<String>,

        /// Build a crate at a specific path
        #[structopt(short = "l", long = "local", conflicts_with_all(&["CRATE_NAME", "CRATE_VERSION"]))]
        local: Option<PathBuf>,
    },

    /// update the currently installed rustup toolchain
    UpdateToolchain {
        /// Update the toolchain only if no toolchain is currently installed
        #[structopt(name = "ONLY_FIRST_TIME", long = "only-first-time")]
        only_first_time: bool,
    },

    /// Adds essential files for the installed version of rustc
    AddEssentialFiles,

    /// Locks cratesfyi daemon to stop building new crates
    Lock,

    /// Unlocks cratesfyi daemon to continue building new crates
    Unlock,

    PrintOptions,
}

impl BuildSubcommand {
    pub fn handle_args(self, mut docbuilder: DocBuilder) {
        match self {
            Self::World => {
                docbuilder.load_cache().expect("Failed to load cache");

                let mut builder = RustwideBuilder::init().unwrap();
                builder
                    .build_world(&mut docbuilder)
                    .expect("Failed to build world");

                docbuilder.save_cache().expect("Failed to save cache");
            }

            Self::Crate {
                crate_name,
                crate_version,
                local,
            } => {
                docbuilder.load_cache().expect("Failed to load cache");
                let mut builder = RustwideBuilder::init().expect("failed to initialize rustwide");

                if let Some(path) = local {
                    builder
                        .build_local_package(&mut docbuilder, &path)
                        .expect("Building documentation failed");
                } else {
                    builder
                        .build_package(
                            &mut docbuilder,
                            &crate_name.unwrap(),
                            &crate_version.unwrap(),
                            None,
                        )
                        .expect("Building documentation failed");
                }

                docbuilder.save_cache().expect("Failed to save cache");
            }

            Self::UpdateToolchain { only_first_time } => {
                if only_first_time {
                    let conn = db::connect_db().unwrap();
                    let res = conn
                        .query("SELECT * FROM config WHERE name = 'rustc_version';", &[])
                        .unwrap();

                    if !res.is_empty() {
                        println!("update-toolchain was already called in the past, exiting");
                        return;
                    }
                }

                let mut builder = RustwideBuilder::init().unwrap();
                builder
                    .update_toolchain()
                    .expect("failed to update toolchain");
            }

            Self::AddEssentialFiles => {
                let mut builder = RustwideBuilder::init().unwrap();
                builder
                    .add_essential_files()
                    .expect("failed to add essential files");
            }

            Self::Lock => docbuilder.lock().expect("Failed to lock"),
            Self::Unlock => docbuilder.unlock().expect("Failed to unlock"),
            Self::PrintOptions => println!("{:?}", docbuilder.options()),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
enum DatabaseSubcommand {
    /// Run database migrations
    Migrate {
        /// The database version to migrate to
        #[structopt(name = "VERSION")]
        version: Option<i64>,
    },

    /// Updates github stats for crates.
    UpdateGithubFields,

    AddDirectory {
        /// Path of file or directory
        #[structopt(name = "DIRECTORY")]
        directory: PathBuf,
        /// Prefix of files in database
        #[structopt(name = "PREFIX", env = "CRATESFYI_PREFIX")]
        prefix: String,
    },

    /// Updates monthly release activity chart
    UpdateReleaseActivity,

    /// Updates search index
    UpdateSearchIndex,

    /// Removes a whole crate from the database
    DeleteCrate {
        /// Name of the crate to delete
        #[structopt(name = "CRATE_NAME")]
        crate_name: String,
    },

    /// Blacklist operations
    Blacklist {
        #[structopt(subcommand)]
        command: BlacklistSubcommand,
    },
}

impl DatabaseSubcommand {
    pub fn handle_args(self) {
        match self {
            Self::Migrate { version } => {
                let conn = connect_db().expect("failed to connect to the database");
                db::migrate(version, &conn).expect("Failed to run database migrations");
            }

            Self::UpdateGithubFields => {
                cratesfyi::utils::github_updater().expect("Failed to update github fields");
            }

            Self::AddDirectory { directory, prefix } => {
                let conn = db::connect_db().expect("failed to connect to the database");
                add_path_into_database(&conn, &prefix, directory)
                    .expect("Failed to add directory into database");
            }

            // FIXME: This is actually util command not database
            Self::UpdateReleaseActivity => cratesfyi::utils::update_release_activity()
                .expect("Failed to update release activity"),

            Self::UpdateSearchIndex => {
                let conn = db::connect_db().expect("failed to connect to the database");
                db::update_search_index(&conn).expect("Failed to update search index");
            }

            Self::DeleteCrate { crate_name } => {
                let conn = db::connect_db().expect("failed to connect to the database");
                db::delete_crate(&conn, &crate_name).expect("failed to delete the crate");
            }

            Self::Blacklist { command } => command.handle_args(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
enum BlacklistSubcommand {
    /// List all crates on the blacklist
    List,

    /// Add a crate to the blacklist
    Add {
        /// Crate name
        #[structopt(name = "CRATE_NAME")]
        crate_name: String,
    },

    /// Remove a crate from the blacklist
    Remove {
        /// Crate name
        #[structopt(name = "CRATE_NAME")]
        crate_name: String,
    },
}

impl BlacklistSubcommand {
    fn handle_args(self) {
        let conn = db::connect_db().expect("failed to connect to the database");

        match self {
            Self::List => {
                let crates =
                    db::blacklist::list_crates(&conn).expect("failed to list crates on blacklist");

                println!("{}", crates.join("\n"));
            }

            Self::Add { crate_name } => db::blacklist::add_crate(&conn, &crate_name)
                .expect("failed to add crate to blacklist"),

            Self::Remove { crate_name } => db::blacklist::remove_crate(&conn, &crate_name)
                .expect("failed to remove crate from blacklist"),
        }
    }
}
<|endoftext|>


<commit_before>#![feature(phase, macro_rules, if_let)]

extern crate serialize;
extern crate libc;
#[phase(plugin, link)] extern crate log;
extern crate regex;
#[phase(plugin)] extern crate regex_macros;
extern crate time;
extern crate debug;

extern crate argparse;
extern crate quire;
#[phase(plugin, link)] extern crate lithos;

use std::rc::Rc;
use std::os::args;
use std::io::stderr;
use std::io::IoError;
use std::io::fs::File;
use std::from_str::FromStr;
use std::io::fs::{readdir};
use std::io::{BufferedReader, MemWriter};
use std::os::{set_exit_status};
use std::default::Default;
use libc::pid_t;
use std::collections::{TreeMap, TreeSet};

use argparse::{ArgumentParser};

use lithos::signal;

#[allow(dead_code, unused_attribute)] mod lithos_tree;
#[allow(dead_code, unused_attribute)] mod lithos_knot;

#[deriving(PartialEq, Eq, PartialOrd, Ord)]
enum LithosInfo {
    Tree(Path),
    Knot(String),
}

#[deriving(Default, PartialEq, Eq, PartialOrd, Ord)]
struct Process {
    parent_id: pid_t,
    pid: pid_t,
    lithos_info: Option<LithosInfo>,
    start_time: uint,
    mem_rss: uint,
    mem_swap: uint,
    threads: uint,
}

fn parse_mem_size(value: &str) -> uint {
    let mut pair = value.as_slice().splitn(1, ' ');
    let val: uint = FromStr::from_str(pair.next().unwrap())
        .expect("Memory should be integer");
    let unit = pair.next().unwrap_or("kB");
    match unit {
        "kB" | "" => val * 1024,
        "MB" => val * 1024 * 1024,
        _ => {
            warn!("Wrong memory unit: {}", unit);
            val * 1024
        }
    }
}

fn read_cmdline(pid: pid_t) -> Result<Vec<String>, IoError> {
    let mut f = try!(File::open(
        &Path::new(format!("/proc/{}/cmdline", pid).as_slice())));
    let mut args: Vec<String> = try!(f.read_to_string())
              .as_slice().split('\0')
              .map(|x| x.to_string())
              .collect();
    args.pop();  // empty arg at the end
    return Ok(args);
}

fn get_tree_info(pid: pid_t) -> Result<LithosInfo, ()> {
    read_cmdline(pid)
    .map_err(|_| debug!("Can't read cmdline for {}", pid))
    .and_then(|args| {
        let mut out = MemWriter::new();
        let mut err = MemWriter::new();
        lithos_tree::Options::parse_specific_args(args, &mut out, &mut err)
            .map_err(|_| debug!("Can't parse lithos_tree cmdline for {}", pid))
    })
    .map(|opt| Tree(opt.config_file))
}

fn get_knot_info(pid: pid_t) -> Result<LithosInfo, ()> {
    read_cmdline(pid)
    .map_err(|_| debug!("Can't read cmdline for {}", pid))
    .and_then(|args| {
        let mut out = MemWriter::new();
        let mut err = MemWriter::new();
        lithos_knot::Options::parse_specific_args(args, &mut out, &mut err)
            .map_err(|_| debug!("Can't parse lithos_knot cmdline for {}", pid))
    })
    .map(|opt| Knot(opt.name))
}

fn read_process(pid: pid_t) -> Result<Process, IoError> {
    let mut f = BufferedReader::new(try!(File::open(
        &Path::new(format!("/proc/{}/status", pid).as_slice()))));
    let mut result: Process = Default::default();
    result.pid = pid;
    for line in f.lines() {
        let line = try!(line);
        let mut pair = line.as_slice().splitn(1, ':');
        let name = pair.next().unwrap().trim();
        let value = pair.next().expect("Colon and value expected").trim();
        match name.as_slice(){
            "PPid" => {
                result.parent_id = FromStr::from_str(value)
                    .expect("Ppid should be integer");
            }
            "VmRSS" => {
                result.mem_rss = parse_mem_size(value);
            }
            "VmSwap" => {
                result.mem_swap = parse_mem_size(value);
            }
            "Threads" => {
                result.threads = FromStr::from_str(value)
                    .expect("Threads should be integer");
            }
            "Name" => {
                match value {
                    "lithos_tree" => {
                        result.lithos_info = get_tree_info(pid).ok();
                    }
                    "lithos_knot" => {
                        result.lithos_info = get_knot_info(pid).ok();
                    }
                    _ => {}
                };
            }
            _ => {}
        }
    }
    let mut f = BufferedReader::new(try!(File::open(
        &Path::new(format!("/proc/{}/stat", pid).as_slice()))));
    let line = try!(f.read_line());
    let starttime_re = regex!(
        // We parse bracketed executable and double-bracketed, still
        // if executable name itself contains bracket we would fail
        // But even if we fail, we get only start_time missed or incorrect
        r"^\d+\s+(?:\([^)]*\)|\(\([^)]*\)\))(?:\s+(\S+)){20}\s+(\d+)\b");
        // TODO(tailhook) we can get executable name from /status and match
        // it here, the executable in /status is escaped!
    result.start_time = starttime_re
        .captures(line.as_slice())
        .map(|c| c.at(1))
        .and_then(|v| FromStr::from_str(v))
        .unwrap_or_else(|| {
            warn!("Error getting start_time for pid {}", pid);
            return 0;
        });

    return Ok(result);
}

fn scan_processes() -> Result<(), IoError>
{
    let mut children = TreeMap::<pid_t, Vec<Rc<Process>>>::new();
    let mut roots = TreeSet::<Rc<Process>>::new();

    for pid in try!(readdir(&Path::new("/proc")))
        .into_iter()
        .filter_map(|p| p.filename_str().and_then(FromStr::from_str))
    {
        match read_process(pid) {
            Ok(prc) => {
                let prc = Rc::new(prc);
                if let Some(Tree(_)) = prc.lithos_info {
                    roots.insert(prc.clone());
                }
                if let Some(vec) = children.find_mut(&prc.parent_id) {
                    vec.push(prc);
                    continue;
                }
                children.insert(prc.parent_id, vec!(prc));
            }
            Err(e) => {
                info!("Error reading pid {}: {}", pid, e);
            }
        }
    }

    for root in roots.iter() {
        if let Some(Tree(ref cfg_file)) = root.lithos_info {
            println!("-+= {} tree ({})", root.pid, cfg_file.display());
            for prc in children.pop(&root.pid).unwrap_or(Vec::new())
                .iter()
            {
                if let Some(Knot(ref name)) = prc.lithos_info {
                    println!(r" \--- {} {}", prc.pid, name);
                }
            }
        }
    }

    return Ok(());
}

fn main() {

    signal::block_all();

    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Displays tree of processes");
        match ap.parse_args() {
            Ok(()) => {}
            Err(x) => {
                set_exit_status(x);
                return;
            }
        }
    }
    match scan_processes() {
        Ok(()) => {
            set_exit_status(0);
        }
        Err(e) => {
            (write!(stderr(), "Fatal error: {}\n", e)).ok();
            set_exit_status(1);
        }
    }
}
<commit_msg>More useful lithos_ps<commit_after>#![feature(phase, macro_rules, if_let, slicing_syntax)]

extern crate serialize;
extern crate libc;
#[phase(plugin, link)] extern crate log;
extern crate regex;
#[phase(plugin)] extern crate regex_macros;
extern crate time;
extern crate debug;

extern crate argparse;
extern crate quire;
#[phase(plugin, link)] extern crate lithos;

use std::rc::Rc;
use std::os::args;
use std::io::stderr;
use std::io::IoError;
use std::io::fs::File;
use std::from_str::FromStr;
use std::io::fs::{readdir};
use std::io::{BufferedReader, MemWriter};
use std::os::{set_exit_status};
use std::default::Default;
use std::collections::{TreeMap, TreeSet};
use time::get_time;
use libc::pid_t;
use libc::consts::os::sysconf::_SC_CLK_TCK;
use libc::funcs::posix88::unistd::sysconf;

use argparse::{ArgumentParser};

use lithos::signal;

#[allow(dead_code, unused_attribute)] mod lithos_tree;
#[allow(dead_code, unused_attribute)] mod lithos_knot;

static mut boot_time: uint = 0;
static mut clock_ticks: uint = 100;

#[deriving(PartialEq, Eq, PartialOrd, Ord)]
enum LithosInfo {
    Tree(Path),
    Knot(String),
}

#[deriving(Default, PartialEq, Eq, PartialOrd, Ord)]
struct Process {
    parent_id: pid_t,
    pid: pid_t,
    lithos_info: Option<LithosInfo>,
    start_time: uint,
    mem_rss: uint,
    mem_swap: uint,
    threads: uint,
    cmdline: Vec<String>,
}

#[deriving(Default)]
struct KnotTotals {
    mem_rss: uint,
    mem_swap: uint,
    processes: uint,
    threads: uint,
}

fn parse_mem_size(value: &str) -> uint {
    let mut pair = value.as_slice().splitn(1, ' ');
    let val: uint = FromStr::from_str(pair.next().unwrap())
        .expect("Memory should be integer");
    let unit = pair.next().unwrap_or("kB");
    match unit {
        "kB" | "" => val * 1024,
        "MB" => val * 1024 * 1024,
        _ => {
            warn!("Wrong memory unit: {}", unit);
            val * 1024
        }
    }
}

fn read_cmdline(pid: pid_t) -> Result<Vec<String>, IoError> {
    let mut f = try!(File::open(
        &Path::new(format!("/proc/{}/cmdline", pid).as_slice())));
    let mut args: Vec<String> = try!(f.read_to_string())
              .as_slice().split('\0')
              .map(|x| x.to_string())
              .collect();
    args.pop();  // empty arg at the end
    return Ok(args);
}

fn get_tree_info(pid: pid_t, cmdline: &Vec<String>) -> Result<LithosInfo, ()> {
    let args = cmdline.clone();
    let mut out = MemWriter::new();
    let mut err = MemWriter::new();
    lithos_tree::Options::parse_specific_args(args, &mut out, &mut err)
        .map(|opt| Tree(opt.config_file))
        .map_err(|_| debug!("Can't parse lithos_tree cmdline for {}", pid))
}

fn get_knot_info(pid: pid_t, cmdline: &Vec<String>) -> Result<LithosInfo, ()> {
    let args = cmdline.clone();
    let mut out = MemWriter::new();
    let mut err = MemWriter::new();
    lithos_knot::Options::parse_specific_args(args, &mut out, &mut err)
        .map(|opt| Knot(opt.name))
        .map_err(|_| debug!("Can't parse lithos_knot cmdline for {}", pid))
}

fn read_process(pid: pid_t) -> Result<Process, IoError> {
    let mut f = BufferedReader::new(try!(File::open(
        &Path::new(format!("/proc/{}/status", pid).as_slice()))));
    let mut result: Process = Default::default();
    result.pid = pid;
    result.cmdline = try!(read_cmdline(pid));
    for line in f.lines() {
        let line = try!(line);
        let mut pair = line.as_slice().splitn(1, ':');
        let name = pair.next().unwrap().trim();
        let value = pair.next().expect("Colon and value expected").trim();
        match name.as_slice(){
            "PPid" => {
                result.parent_id = FromStr::from_str(value)
                    .expect("Ppid should be integer");
            }
            "VmRSS" => {
                result.mem_rss = parse_mem_size(value);
            }
            "VmSwap" => {
                result.mem_swap = parse_mem_size(value);
            }
            "Threads" => {
                result.threads = FromStr::from_str(value)
                    .expect("Threads should be integer");
            }
            "Name" => {
                match value {
                    "lithos_tree" => {
                        result.lithos_info = get_tree_info(
                            pid, &result.cmdline).ok();
                    }
                    "lithos_knot" => {
                        result.lithos_info = get_knot_info(
                            pid, &result.cmdline).ok();
                    }
                    _ => {}
                };
            }
            _ => {}
        }
    }
    let mut f = BufferedReader::new(try!(File::open(
        &Path::new(format!("/proc/{}/stat", pid).as_slice()))));
    let line = try!(f.read_line());
    let starttime_re = regex!(
        // We parse bracketed executable and double-bracketed, still
        // if executable name itself contains bracket we would fail
        // But even if we fail, we get only start_time missed or incorrect
        r"^\d+\s+(?:\([^)]*\)|\(\([^)]*\)\))(?:\s+(\S+)){20}\s+(\d+)\b");
        // TODO(tailhook) we can get executable name from /status and match
        // it here, the executable in /status is escaped!
    result.start_time = starttime_re
        .captures(line.as_slice())
        .map(|c| c.at(1))
        .and_then(|v| FromStr::from_str(v))
        .unwrap_or_else(|| {
            warn!("Error getting start_time for pid {}", pid);
            return 0;
        });

    return Ok(result);
}

impl KnotTotals {
    fn _add_process(&mut self, prc: &Process,
        children: &TreeMap<pid_t, Vec<Rc<Process>>>)
    {
        self.mem_rss += prc.mem_rss;
        self.mem_swap += prc.mem_swap;
        self.processes += 1;
        self.threads += prc.threads;
        if let Some(ref lst) = children.find(&prc.pid) {
            for prc in lst.iter() {
                self._add_process(&**prc, children);
            }
        }
    }
}

fn format_uptime(start_ticks: uint) -> String {
    let start_time = unsafe { boot_time  + (start_ticks / clock_ticks) };
    let uptime = get_time().sec as uint - start_time;
    if uptime < 60 {
        format!("{}s", uptime)
    } else if uptime < 3600 {
        format!("{}m{}s", uptime / 60, uptime % 60)
    } else if uptime < 86400 {
        format!("{}h{}m{}s", uptime / 3600, (uptime / 60) % 60, uptime % 60)
    } else if uptime < 3*86400 {
        format!("{}d{}h{}m{}s", uptime / 86400,
            (uptime / 3600) % 24, (uptime / 60) % 60, uptime % 60)
    } else {
        format!("{}days", uptime / 86400)
    }
}

fn format_memory(mem: uint) -> String {
    if mem < (1 << 10) {
        format!("{}B", mem)
    } else if mem < (1 << 20) {
        format!("{:.1f}kiB", mem as f64/(1024_f64))
    } else if mem < (1 << 30) {
        format!("{:.1f}MiB", mem as f64/(1_048_576_f64))
    } else {
        format!("{:.1f}GiB", mem as f64/(1_048_576_f64 * 1024_f64))
    }
}

fn scan_processes() -> Result<(), IoError>
{
    let mut children = TreeMap::<pid_t, Vec<Rc<Process>>>::new();
    let mut roots = TreeSet::<Rc<Process>>::new();

    for pid in try!(readdir(&Path::new("/proc")))
        .into_iter()
        .filter_map(|p| p.filename_str().and_then(FromStr::from_str))
    {
        match read_process(pid) {
            Ok(prc) => {
                let prc = Rc::new(prc);
                if let Some(Tree(_)) = prc.lithos_info {
                    roots.insert(prc.clone());
                }
                if let Some(vec) = children.find_mut(&prc.parent_id) {
                    vec.push(prc);
                    continue;
                }
                children.insert(prc.parent_id, vec!(prc));
            }
            Err(e) => {
                info!("Error reading pid {}: {}", pid, e);
            }
        }
    }

    for root in roots.iter() {
        if let Some(Tree(ref cfg_file)) = root.lithos_info {
            println!("-+= {} tree ({})", root.pid, cfg_file.display());
            for prc in children.find(&root.pid).unwrap_or(&Vec::new())
                .iter()
            {
                if let Some(Knot(ref name)) = prc.lithos_info {
                    if let Some(knot_children) = children.find(&prc.pid) {
                        if knot_children.len() == 1 {
                            let child = &knot_children[0];
                            let mut info: KnotTotals = Default::default();
                            info._add_process(&**child, &children);
                            println!(r" \--- {} {} {} [{}/{}] {}",
                                child.pid, name,
                                format_uptime(child.start_time),
                                info.processes,
                                info.threads,
                                format_memory(info.mem_rss + info.mem_swap));
                        } else {
                            println!(r" \-+- ({}) {} [multiple]",
                                     prc.pid, name);
                            for child in knot_children.iter() {
                                let mut info: KnotTotals = Default::default();
                                info._add_process(&**child, &children);
                                println!(r"   \--- {} {} {} [{}/{}] {}",
                                    child.pid, name,
                                    format_uptime(child.start_time),
                                    info.processes,
                                    info.threads,
                                    format_memory(info.mem_rss+info.mem_swap),
                                    );
                            }
                        }
                    } else {
                            println!(r" \--- ({}) {} [failing]",
                                     prc.pid, name);
                    }
                }
            }
        }
    }

    return Ok(());
}

fn read_global_consts() {
    unsafe {
        clock_ticks = sysconf(_SC_CLK_TCK) as uint;
        boot_time =
            BufferedReader::new(
                File::open(&Path::new("/proc/stat"))
                .ok().expect("Can't read /proc/stat"))
            .lines()
            .map(|line| line.ok().expect("Can't read /proc/stat"))
            .filter(|line| line.as_slice().starts_with("btime "))
            .next()
            .and_then(|line| FromStr::from_str(line.as_slice()[5..].trim()))
            .expect("No boot time in /proc/stat");
    }
}

fn main() {

    signal::block_all();

    read_global_consts();

    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Displays tree of processes");
        match ap.parse_args() {
            Ok(()) => {}
            Err(x) => {
                set_exit_status(x);
                return;
            }
        }
    }
    match scan_processes() {
        Ok(()) => {
            set_exit_status(0);
        }
        Err(e) => {
            (write!(stderr(), "Fatal error: {}\n", e)).ok();
            set_exit_status(1);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

#[macro_use]
extern crate log;
extern crate env_logger;
extern crate osmpbfreader;
extern crate rustc_serialize;
extern crate docopt;
extern crate mimirsbrunn;
extern crate rs_es;

use std::collections::HashSet;
use std::collections::BTreeMap;
use osmpbfreader::{OsmObj, OsmId};
use mimirsbrunn::rubber::Rubber;

pub type AdminsMap = BTreeMap<OsmId, mimirsbrunn::Admin>;

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_input: String,
    flag_level: Vec<u32>,
    flag_connection_string: String,
}

static USAGE: &'static str = "
Usage:
    osm2mimir --help
    osm2mimir --input=<file> [--connection-string=<connection-string>] --level=<level>...

Options:
    -h, --help            Show this message.
    -i, --input=<file>    OSM PBF file.
    -l, --level=<level>   Admin levels to keep.
    -c, --connection-string=<connection-string>
                          Elasticsearch parameters, [default: http://localhost:9200/munin]
";

fn get_osm_id(obj: &OsmObj) -> OsmId {
    match *obj {
        OsmObj::Node(ref node) => OsmId::Node(node.id),
        OsmObj::Way(ref way) => OsmId::Way(way.id),
        OsmObj::Relation(ref rel) => OsmId::Relation(rel.id),
    }
}

fn update_coordinates(filename: &String, admins: &mut AdminsMap) {
    if admins.is_empty() {
        return;
    }
    // load coord for administratives regions
    let path = std::path::Path::new(&filename);
    let r = std::fs::File::open(&path).unwrap();
    let mut pbf = osmpbfreader::OsmPbfReader::new(r);
    for obj in pbf.iter() {
        if let osmpbfreader::OsmObj::Node(ref node) = obj {
            let mut adm = match admins.get_mut(&get_osm_id(&obj)) {
                Some(val) => val,
                None => continue,
            };
            adm.coord.lat = node.lat;
            adm.coord.lon = node.lon;
        }
    }
}

fn administartive_regions(filename: &String, levels: &HashSet<u32>) -> AdminsMap {
    let mut administrative_regions = AdminsMap::new();
    let path = std::path::Path::new(&filename);
    let r = std::fs::File::open(&path).unwrap();
    let mut pbf = osmpbfreader::OsmPbfReader::new(r);
    // load administratives regions
    for obj in pbf.iter() {
        if let osmpbfreader::OsmObj::Relation(relation) = obj {
            // not administartive region
            if !relation.tags
                        .get("boundary")
                        .map(|s| s == "administrative")
                        .unwrap_or(false) {
                continue;
            }
            let level = relation.tags
                                .get("admin_level")
                                .and_then(|s| s.parse().ok());
            let level = match level {
                None => {
                    info!("invalid admin_level for relation {}: admin_level {:?}",
                          relation.id,
                          relation.tags.get("admin_level"));
                    continue;
                }
                Some(ref l) if !levels.contains(&l) => continue,
                Some(l) => l,
            };
            // administrative region with name ?
            let name = match relation.tags.get("name") {
                Some(val) => val,
                None => {
                    info!("adminstrative region without name for relation {}:  admin_level {} \
                           ignored.",
                          relation.id,
                          level);
                    continue;
                }
            };
            // admininstrative region without coordinates
            let admin_centre = match relation.refs.iter().find(|rf| rf.role == "admin_centre") {
                Some(val) => val.member,
                None => {
                    info!("adminstrative region [{}] without coordinates for relation {}.",
                          name,
                          relation.id);
                    continue;
                }
            };

            let admin_id = match relation.tags.get("ref:INSEE") {
                Some(val) => format!("admin:fr:{}", val),
                None => format!("admin:osm:{}", relation.id),
            };
            let zip_code = match relation.tags.get("addr:postcode") {
                Some(val) => &val[..],
                None => "",
            };
            let admin = mimirsbrunn::Admin {
                id: admin_id,
                level: level,
                name: name.to_string(),
                zip_code: zip_code.to_string(),
                // TODO weight value ?
                weight: 1,
                coord: mimirsbrunn::Coord {
                    lat: 0.0,
                    lon: 0.0,
                },
            };
            administrative_regions.insert(admin_centre, admin);
        }
    }
    return administrative_regions;
}

fn index_osm(es_cnx_string: &str, admins: &AdminsMap) -> Result<u32, rs_es::error::EsError> {
    let mut rubber = Rubber::new(es_cnx_string);
    info!("purge and create Munin...");
    try!(rubber.create_index());
    info!("Munin purged and created.");
    rubber.bulk_index(admins.values())
}

fn main() {
    env_logger::init().unwrap();
    debug!("importing adminstrative region into Mimir");
    let args: Args = docopt::Docopt::new(USAGE)
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());

    let levels = args.flag_level.iter().cloned().collect();
    let mut res = administartive_regions(&args.flag_input, &levels);
    update_coordinates(&args.flag_input, &mut res);
    match index_osm(&args.flag_connection_string, &res) {
        Err(e) => panic!("failed to index osm because: {}", e),
        Ok(nb) => info!("Adminstrative regions : {}", nb),
    }
}
<commit_msg>Use OsmObj::id()<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

#[macro_use]
extern crate log;
extern crate env_logger;
extern crate osmpbfreader;
extern crate rustc_serialize;
extern crate docopt;
extern crate mimirsbrunn;
extern crate rs_es;

use std::collections::HashSet;
use std::collections::BTreeMap;
use osmpbfreader::OsmId;
use mimirsbrunn::rubber::Rubber;

pub type AdminsMap = BTreeMap<OsmId, mimirsbrunn::Admin>;

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_input: String,
    flag_level: Vec<u32>,
    flag_connection_string: String,
}

static USAGE: &'static str = "
Usage:
    osm2mimir --help
    osm2mimir --input=<file> [--connection-string=<connection-string>] --level=<level>...

Options:
    -h, --help            Show this message.
    -i, --input=<file>    OSM PBF file.
    -l, --level=<level>   Admin levels to keep.
    -c, --connection-string=<connection-string>
                          Elasticsearch parameters, [default: http://localhost:9200/munin]
";

fn update_coordinates(filename: &String, admins: &mut AdminsMap) {
    if admins.is_empty() {
        return;
    }
    // load coord for administratives regions
    let path = std::path::Path::new(&filename);
    let r = std::fs::File::open(&path).unwrap();
    let mut pbf = osmpbfreader::OsmPbfReader::new(r);
    for obj in pbf.iter() {
        if let osmpbfreader::OsmObj::Node(ref node) = obj {
            let mut adm = match admins.get_mut(&obj.id()) {
                Some(val) => val,
                None => continue,
            };
            adm.coord.lat = node.lat;
            adm.coord.lon = node.lon;
        }
    }
}

fn administartive_regions(filename: &String, levels: &HashSet<u32>) -> AdminsMap {
    let mut administrative_regions = AdminsMap::new();
    let path = std::path::Path::new(&filename);
    let r = std::fs::File::open(&path).unwrap();
    let mut pbf = osmpbfreader::OsmPbfReader::new(r);
    // load administratives regions
    for obj in pbf.iter() {
        if let osmpbfreader::OsmObj::Relation(relation) = obj {
            // not administartive region
            if !relation.tags
                        .get("boundary")
                        .map(|s| s == "administrative")
                        .unwrap_or(false) {
                continue;
            }
            let level = relation.tags
                                .get("admin_level")
                                .and_then(|s| s.parse().ok());
            let level = match level {
                None => {
                    info!("invalid admin_level for relation {}: admin_level {:?}",
                          relation.id,
                          relation.tags.get("admin_level"));
                    continue;
                }
                Some(ref l) if !levels.contains(&l) => continue,
                Some(l) => l,
            };
            // administrative region with name ?
            let name = match relation.tags.get("name") {
                Some(val) => val,
                None => {
                    info!("adminstrative region without name for relation {}:  admin_level {} \
                           ignored.",
                          relation.id,
                          level);
                    continue;
                }
            };
            // admininstrative region without coordinates
            let admin_centre = match relation.refs.iter().find(|rf| rf.role == "admin_centre") {
                Some(val) => val.member,
                None => {
                    info!("adminstrative region [{}] without coordinates for relation {}.",
                          name,
                          relation.id);
                    continue;
                }
            };

            let admin_id = match relation.tags.get("ref:INSEE") {
                Some(val) => format!("admin:fr:{}", val),
                None => format!("admin:osm:{}", relation.id),
            };
            let zip_code = match relation.tags.get("addr:postcode") {
                Some(val) => &val[..],
                None => "",
            };
            let admin = mimirsbrunn::Admin {
                id: admin_id,
                level: level,
                name: name.to_string(),
                zip_code: zip_code.to_string(),
                // TODO weight value ?
                weight: 1,
                coord: mimirsbrunn::Coord {
                    lat: 0.0,
                    lon: 0.0,
                },
            };
            administrative_regions.insert(admin_centre, admin);
        }
    }
    return administrative_regions;
}

fn index_osm(es_cnx_string: &str, admins: &AdminsMap) -> Result<u32, rs_es::error::EsError> {
    let mut rubber = Rubber::new(es_cnx_string);
    info!("purge and create Munin...");
    try!(rubber.create_index());
    info!("Munin purged and created.");
    rubber.bulk_index(admins.values())
}

fn main() {
    env_logger::init().unwrap();
    debug!("importing adminstrative region into Mimir");
    let args: Args = docopt::Docopt::new(USAGE)
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());

    let levels = args.flag_level.iter().cloned().collect();
    let mut res = administartive_regions(&args.flag_input, &levels);
    update_coordinates(&args.flag_input, &mut res);
    match index_osm(&args.flag_connection_string, &res) {
        Err(e) => panic!("failed to index osm because: {}", e),
        Ok(nb) => info!("Adminstrative regions : {}", nb),
    }
}
<|endoftext|>


<commit_before>extern crate luthor;

use std::rc::Rc;
use std::cell::RefCell;
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::path::PathBuf;
use super::GapBuffer;
use super::gap_buffer;
use super::Position;
use super::Range;
use super::Cursor;
use super::cursor;
use super::type_detection;
use self::luthor::token::Token;
use self::luthor::lexers;

/// A UTF-8 buffer with bounds-checked cursor management and persistence.
pub struct Buffer {
    data: Rc<RefCell<GapBuffer>>,
    lexer: fn(&str) -> Vec<Token>,
    pub path: Option<PathBuf>,
    pub cursor: Cursor,
    data_cache: Option<String>,
    token_cache: Option<Vec<Token>>,
}

impl Buffer {
    /// Returns the contents of the buffer as a string. Caches
    /// this string representation to make subsequent requests
    /// to an unchanged buffer as fast as possible.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// assert_eq!(buffer.data(), "scribe");
    /// ```
    pub fn data(&mut self) -> String {
        match self.data_cache {
            Some(ref cache) => cache.clone(), // Cache hit; return a copy.
            None => {
                // Cache miss; update the cache w/ fresh data and return a copy.
                self.data_cache = Some(self.data.borrow().to_string());
                self.data_cache.clone().unwrap()
            }
        }
    }

    /// Writes the contents of the buffer to its path.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::path::PathBuf;
    /// # use std::path::Path;
    /// # use std::fs::File;
    /// # use std::io::Read;
    ///
    /// // Set up a buffer and point it to a path.
    /// let mut buffer = scribe::buffer::new();
    /// let write_path = PathBuf::from("my_doc");
    /// buffer.path = Some(write_path.clone());
    ///
    /// // Put some data into the buffer and save it.
    /// buffer.insert("scribe");
    /// buffer.save();
    ///
    /// # let mut saved_data = String::new();
    /// # File::open(Path::new("my_doc")).unwrap().
    /// #   read_to_string(&mut saved_data).unwrap();
    /// # assert_eq!(saved_data, "scribe");
    ///
    /// # std::fs::remove_file(&write_path);
    /// ```
    pub fn save(&mut self) -> Option<io::Error> {
        let path = match self.path.clone() {
            Some(p) => p,
            None => PathBuf::new(),
        };

        // Try to open and write to the file, returning any errors encountered.
        let mut file = match File::create(&path) {
            Ok(f) => f,
            Err(error) => return Some(error),
        };

        // We use to_string here because we don't want to write the gap contents.
        match file.write_all(self.data().to_string().as_bytes()) {
            Ok(_) => (),
            Err(error) => return Some(error),
        }

        return None
    }

    /// Inserts `data` into the buffer at the cursor position.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// assert_eq!(buffer.data(), "scribe");
    /// ```
    pub fn insert(&mut self, data: &str) {
        // Insert the data at the cursor
        self.data.borrow_mut().insert(data, &self.cursor);

        // Caches are invalid as the buffer has changed.
        self.clear_caches();
    }

    /// Deletes a character at the cursor position. If at the end
    /// of the current line, it'll try to delete a newline character
    /// (joining the lines), succeeding if there's a line below.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// buffer.delete();
    /// assert_eq!(buffer.data(), "cribe");
    /// ```
    pub fn delete(&mut self) {
        // We need to specify a range to delete, so start at
        // the current offset and delete the character to the right.
        let mut end = Position{ line: self.cursor.line, offset: self.cursor.offset + 1 };

        // If there isn't a character to the right,
        // delete the newline by jumping to the start
        // of the next line. If it doesn't exist, that's okay;
        // these values are bounds-checked by delete() anyway.
        if !self.data.borrow().in_bounds(&end) {
            end.line += 1;
            end.offset = 0;
        }

        // Delete the data.
        self.data.borrow_mut().delete(&Range{ start: *self.cursor, end: end});

        // Caches are invalid as the buffer has changed.
        self.clear_caches();
    }

    /// Produces a set of tokens based on the buffer data
    /// suitable for colorized display, using a lexer for the
    /// buffer data's language and/or format. Caches this
    /// lexed representation to make subsequent requests
    /// to an unchanged buffer as fast as possible.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe data");
    ///
    /// // Build the buffer data string back by combining its token lexemes.
    /// let mut data = String::new();
    /// for token in buffer.tokens().iter() {
    ///     data.push_str(&token.lexeme);
    /// }
    /// assert_eq!(data, "scribe data");
    /// ```
    pub fn tokens(&mut self) -> Vec<Token> {
        match self.token_cache {
            Some(ref cache) => cache.clone(), // Cache hit; return a copy.
            None => {
                // Cache miss; update the cache w/ fresh tokens and return a copy.
                self.token_cache = Some((self.lexer)(&self.data()));
                self.token_cache.clone().unwrap()
            }
        }
    }

    /// Returns the file name portion of the buffer's path, if
    /// the path is set and its file name is a valid UTF-8 sequence.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::path::PathBuf;
    ///
    /// let file_path = PathBuf::from("tests/sample/file");
    /// let buffer = scribe::buffer::from_file(file_path).unwrap();
    /// assert_eq!(buffer.file_name().unwrap(), "file");
    /// ```
    pub fn file_name(&self) -> Option<String> {
        match self.path {
            Some(ref path) => {
                match path.file_name() {
                    Some(file_name) => {
                        match file_name.to_str() {
                            Some(utf8_file_name) => Some(utf8_file_name.to_string()),
                            None => None,
                        }
                    },
                    None => None,
                }
            },
            None => None,
        }
    }

    /// Called when caches are invalidated via buffer modifications.
    fn clear_caches(&mut self) {
        self.data_cache = None;
        self.token_cache = None;
    }
}

/// Creates a new empty buffer. The buffer's cursor is set to the beginning of the buffer.
///
/// # Examples
///
/// ```
/// let buffer = scribe::buffer::new();
/// # assert_eq!(buffer.cursor.line, 0);
/// # assert_eq!(buffer.cursor.offset, 0);
/// ```
pub fn new() -> Buffer {
    let data = Rc::new(RefCell::new(gap_buffer::new(String::new())));
    let cursor = cursor::new(data.clone(), 0, 0);

    Buffer{
        data: data.clone(),
        path: None,
        cursor: cursor,
        lexer: lexers::default::lex as fn(&str) -> Vec<Token>,
        data_cache: None,
        token_cache: None,
    }
}

/// Creates a new buffer by reading the UTF-8 interpreted file contents of the specified path.
/// The buffer's cursor is set to the beginning of the buffer. The buffer data's type will be
/// inferred based on its extension, and an appropriate lexer will be used, if available (see
/// tokens method for further information on why this happens).
///
/// # Examples
///
/// ```
/// use std::path::PathBuf;
///
/// let file_path = PathBuf::from("tests/sample/file");
/// let mut buffer = scribe::buffer::from_file(file_path).unwrap();
/// assert_eq!(buffer.data(), "it works!\n");
/// # assert_eq!(buffer.cursor.line, 0);
/// # assert_eq!(buffer.cursor.offset, 0);
/// ```
pub fn from_file(path: PathBuf) -> io::Result<Buffer> {
    // Try to open and read the file, returning any errors encountered.
    let mut file = match File::open(path.clone()) {
        Ok(f) => f,
        Err(error) => return Err(error),
    };
    let mut data = String::new();
    match file.read_to_string(&mut data) {
        Ok(_) => (),
        Err(error) => return Err(error),
    };

    let data = Rc::new(RefCell::new(gap_buffer::new(data)));
    let cursor = cursor::new(data.clone(), 0, 0);

    // Detect the file type and use its corresponding lexer, if available.
    let lexer = match type_detection::from_path(&path) {
        Some(type_detection::Type::JSON) => lexers::json::lex as fn(&str) -> Vec<Token>,
        Some(type_detection::Type::XML) => lexers::xml::lex as fn(&str) -> Vec<Token>,
        Some(type_detection::Type::Ruby) => lexers::ruby::lex as fn(&str) -> Vec<Token>,
        _ => lexers::default::lex as fn(&str) -> Vec<Token>,
    };

    // Create a new buffer using the loaded data, path, and other defaults.
    Ok(
        Buffer{
            data: data.clone(),
            path: Some(path),
            cursor: cursor,
            lexer: lexer,
            data_cache: None,
            token_cache: None,
        }
    )
}

#[cfg(test)]
mod tests {
    use super::new;
    use super::luthor::token::{Token, Category};

    #[test]
    fn tokens_returns_result_of_lexer() {
        let mut buffer = new();
        buffer.insert("scribe data");
        let expected_tokens = vec![
            Token{ lexeme: "scribe".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "data".to_string(), category: Category::Text },
        ];
        assert_eq!(buffer.tokens(), expected_tokens);
    }

    #[test]
    fn delete_joins_lines_when_invoked_at_end_of_line() {
        let mut buffer = new();
        buffer.insert("scribe\n library");
        buffer.cursor.move_to_end_of_line();
        buffer.delete();
        assert_eq!(buffer.data(), "scribe library");
    }

    #[test]
    fn delete_does_nothing_when_invoked_at_the_end_of_the_document() {
        let mut buffer = new();
        buffer.insert("scribe\n library");
        buffer.cursor.move_down();
        buffer.cursor.move_to_end_of_line();
        buffer.delete();
        assert_eq!(buffer.data(), "scribe\n library");
    }

    #[test]
    fn insert_clears_data_and_token_caches() {
        let mut buffer = new();
        buffer.insert("scribe");

        // Trigger data and token cache storage.
        assert_eq!("scribe", buffer.data());
        assert_eq!(
            vec![Token{ lexeme: "scribe".to_string(), category: Category::Text }],
            buffer.tokens()
        );

        // Change the buffer contents using delete.
        buffer.insert("test_");

        // Ensure the cache has been busted.
        assert_eq!("test_scribe", buffer.data());
        assert_eq!(
            vec![Token{ lexeme: "test_scribe".to_string(), category: Category::Text }],
            buffer.tokens()
        );
    }

    #[test]
    fn delete_clears_data_and_token_caches() {
        let mut buffer = new();
        buffer.insert("scribe");

        // Trigger data and token cache storage.
        assert_eq!("scribe", buffer.data());
        assert_eq!(
            vec![Token{ lexeme: "scribe".to_string(), category: Category::Text }],
            buffer.tokens()
        );

        // Change the buffer contents using delete.
        buffer.delete();

        // Ensure the cache has been busted.
        assert_eq!("cribe", buffer.data());
        assert_eq!(
            vec![Token{ lexeme: "cribe".to_string(), category: Category::Text }],
            buffer.tokens()
        );
    }
}
<commit_msg>Update buffer type to use insert operation.<commit_after>extern crate luthor;

use std::rc::Rc;
use std::cell::RefCell;
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::path::PathBuf;
use super::{GapBuffer, gap_buffer, Position, Range, Cursor, cursor, type_detection, operations};
use super::operations::Operation;
use self::luthor::token::Token;
use self::luthor::lexers;

/// A UTF-8 buffer with bounds-checked cursor management and persistence.
pub struct Buffer {
    data: Rc<RefCell<GapBuffer>>,
    lexer: fn(&str) -> Vec<Token>,
    pub path: Option<PathBuf>,
    pub cursor: Cursor,
    data_cache: Option<String>,
    token_cache: Option<Vec<Token>>,
}

impl Buffer {
    /// Returns the contents of the buffer as a string. Caches
    /// this string representation to make subsequent requests
    /// to an unchanged buffer as fast as possible.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// assert_eq!(buffer.data(), "scribe");
    /// ```
    pub fn data(&mut self) -> String {
        match self.data_cache {
            Some(ref cache) => cache.clone(), // Cache hit; return a copy.
            None => {
                // Cache miss; update the cache w/ fresh data and return a copy.
                self.data_cache = Some(self.data.borrow().to_string());
                self.data_cache.clone().unwrap()
            }
        }
    }

    /// Writes the contents of the buffer to its path.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::path::PathBuf;
    /// # use std::path::Path;
    /// # use std::fs::File;
    /// # use std::io::Read;
    ///
    /// // Set up a buffer and point it to a path.
    /// let mut buffer = scribe::buffer::new();
    /// let write_path = PathBuf::from("my_doc");
    /// buffer.path = Some(write_path.clone());
    ///
    /// // Put some data into the buffer and save it.
    /// buffer.insert("scribe");
    /// buffer.save();
    ///
    /// # let mut saved_data = String::new();
    /// # File::open(Path::new("my_doc")).unwrap().
    /// #   read_to_string(&mut saved_data).unwrap();
    /// # assert_eq!(saved_data, "scribe");
    ///
    /// # std::fs::remove_file(&write_path);
    /// ```
    pub fn save(&mut self) -> Option<io::Error> {
        let path = match self.path.clone() {
            Some(p) => p,
            None => PathBuf::new(),
        };

        // Try to open and write to the file, returning any errors encountered.
        let mut file = match File::create(&path) {
            Ok(f) => f,
            Err(error) => return Some(error),
        };

        // We use to_string here because we don't want to write the gap contents.
        match file.write_all(self.data().to_string().as_bytes()) {
            Ok(_) => (),
            Err(error) => return Some(error),
        }

        return None
    }

    /// Inserts `data` into the buffer at the cursor position.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// assert_eq!(buffer.data(), "scribe");
    /// ```
    pub fn insert(&mut self, data: &str) {
        // Build and run an insert operation.
        let mut op = operations::insert::new(data.to_string(), self.cursor.position.clone());
        op.run(&mut *self.data.borrow_mut());

        // Caches are invalid as the buffer has changed.
        self.clear_caches();
    }

    /// Deletes a character at the cursor position. If at the end
    /// of the current line, it'll try to delete a newline character
    /// (joining the lines), succeeding if there's a line below.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// buffer.delete();
    /// assert_eq!(buffer.data(), "cribe");
    /// ```
    pub fn delete(&mut self) {
        // We need to specify a range to delete, so start at
        // the current offset and delete the character to the right.
        let mut end = Position{ line: self.cursor.line, offset: self.cursor.offset + 1 };

        // If there isn't a character to the right,
        // delete the newline by jumping to the start
        // of the next line. If it doesn't exist, that's okay;
        // these values are bounds-checked by delete() anyway.
        if !self.data.borrow().in_bounds(&end) {
            end.line += 1;
            end.offset = 0;
        }

        // Delete the data.
        self.data.borrow_mut().delete(&Range{ start: *self.cursor, end: end});

        // Caches are invalid as the buffer has changed.
        self.clear_caches();
    }

    /// Produces a set of tokens based on the buffer data
    /// suitable for colorized display, using a lexer for the
    /// buffer data's language and/or format. Caches this
    /// lexed representation to make subsequent requests
    /// to an unchanged buffer as fast as possible.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe data");
    ///
    /// // Build the buffer data string back by combining its token lexemes.
    /// let mut data = String::new();
    /// for token in buffer.tokens().iter() {
    ///     data.push_str(&token.lexeme);
    /// }
    /// assert_eq!(data, "scribe data");
    /// ```
    pub fn tokens(&mut self) -> Vec<Token> {
        match self.token_cache {
            Some(ref cache) => cache.clone(), // Cache hit; return a copy.
            None => {
                // Cache miss; update the cache w/ fresh tokens and return a copy.
                self.token_cache = Some((self.lexer)(&self.data()));
                self.token_cache.clone().unwrap()
            }
        }
    }

    /// Returns the file name portion of the buffer's path, if
    /// the path is set and its file name is a valid UTF-8 sequence.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::path::PathBuf;
    ///
    /// let file_path = PathBuf::from("tests/sample/file");
    /// let buffer = scribe::buffer::from_file(file_path).unwrap();
    /// assert_eq!(buffer.file_name().unwrap(), "file");
    /// ```
    pub fn file_name(&self) -> Option<String> {
        match self.path {
            Some(ref path) => {
                match path.file_name() {
                    Some(file_name) => {
                        match file_name.to_str() {
                            Some(utf8_file_name) => Some(utf8_file_name.to_string()),
                            None => None,
                        }
                    },
                    None => None,
                }
            },
            None => None,
        }
    }

    /// Called when caches are invalidated via buffer modifications.
    fn clear_caches(&mut self) {
        self.data_cache = None;
        self.token_cache = None;
    }
}

/// Creates a new empty buffer. The buffer's cursor is set to the beginning of the buffer.
///
/// # Examples
///
/// ```
/// let buffer = scribe::buffer::new();
/// # assert_eq!(buffer.cursor.line, 0);
/// # assert_eq!(buffer.cursor.offset, 0);
/// ```
pub fn new() -> Buffer {
    let data = Rc::new(RefCell::new(gap_buffer::new(String::new())));
    let cursor = cursor::new(data.clone(), 0, 0);

    Buffer{
        data: data.clone(),
        path: None,
        cursor: cursor,
        lexer: lexers::default::lex as fn(&str) -> Vec<Token>,
        data_cache: None,
        token_cache: None,
    }
}

/// Creates a new buffer by reading the UTF-8 interpreted file contents of the specified path.
/// The buffer's cursor is set to the beginning of the buffer. The buffer data's type will be
/// inferred based on its extension, and an appropriate lexer will be used, if available (see
/// tokens method for further information on why this happens).
///
/// # Examples
///
/// ```
/// use std::path::PathBuf;
///
/// let file_path = PathBuf::from("tests/sample/file");
/// let mut buffer = scribe::buffer::from_file(file_path).unwrap();
/// assert_eq!(buffer.data(), "it works!\n");
/// # assert_eq!(buffer.cursor.line, 0);
/// # assert_eq!(buffer.cursor.offset, 0);
/// ```
pub fn from_file(path: PathBuf) -> io::Result<Buffer> {
    // Try to open and read the file, returning any errors encountered.
    let mut file = match File::open(path.clone()) {
        Ok(f) => f,
        Err(error) => return Err(error),
    };
    let mut data = String::new();
    match file.read_to_string(&mut data) {
        Ok(_) => (),
        Err(error) => return Err(error),
    };

    let data = Rc::new(RefCell::new(gap_buffer::new(data)));
    let cursor = cursor::new(data.clone(), 0, 0);

    // Detect the file type and use its corresponding lexer, if available.
    let lexer = match type_detection::from_path(&path) {
        Some(type_detection::Type::JSON) => lexers::json::lex as fn(&str) -> Vec<Token>,
        Some(type_detection::Type::XML) => lexers::xml::lex as fn(&str) -> Vec<Token>,
        Some(type_detection::Type::Ruby) => lexers::ruby::lex as fn(&str) -> Vec<Token>,
        _ => lexers::default::lex as fn(&str) -> Vec<Token>,
    };

    // Create a new buffer using the loaded data, path, and other defaults.
    Ok(
        Buffer{
            data: data.clone(),
            path: Some(path),
            cursor: cursor,
            lexer: lexer,
            data_cache: None,
            token_cache: None,
        }
    )
}

#[cfg(test)]
mod tests {
    use super::new;
    use super::luthor::token::{Token, Category};

    #[test]
    fn tokens_returns_result_of_lexer() {
        let mut buffer = new();
        buffer.insert("scribe data");
        let expected_tokens = vec![
            Token{ lexeme: "scribe".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "data".to_string(), category: Category::Text },
        ];
        assert_eq!(buffer.tokens(), expected_tokens);
    }

    #[test]
    fn delete_joins_lines_when_invoked_at_end_of_line() {
        let mut buffer = new();
        buffer.insert("scribe\n library");
        buffer.cursor.move_to_end_of_line();
        buffer.delete();
        assert_eq!(buffer.data(), "scribe library");
    }

    #[test]
    fn delete_does_nothing_when_invoked_at_the_end_of_the_document() {
        let mut buffer = new();
        buffer.insert("scribe\n library");
        buffer.cursor.move_down();
        buffer.cursor.move_to_end_of_line();
        buffer.delete();
        assert_eq!(buffer.data(), "scribe\n library");
    }

    #[test]
    fn insert_clears_data_and_token_caches() {
        let mut buffer = new();
        buffer.insert("scribe");

        // Trigger data and token cache storage.
        assert_eq!("scribe", buffer.data());
        assert_eq!(
            vec![Token{ lexeme: "scribe".to_string(), category: Category::Text }],
            buffer.tokens()
        );

        // Change the buffer contents using delete.
        buffer.insert("test_");

        // Ensure the cache has been busted.
        assert_eq!("test_scribe", buffer.data());
        assert_eq!(
            vec![Token{ lexeme: "test_scribe".to_string(), category: Category::Text }],
            buffer.tokens()
        );
    }

    #[test]
    fn delete_clears_data_and_token_caches() {
        let mut buffer = new();
        buffer.insert("scribe");

        // Trigger data and token cache storage.
        assert_eq!("scribe", buffer.data());
        assert_eq!(
            vec![Token{ lexeme: "scribe".to_string(), category: Category::Text }],
            buffer.tokens()
        );

        // Change the buffer contents using delete.
        buffer.delete();

        // Ensure the cache has been busted.
        assert_eq!("cribe", buffer.data());
        assert_eq!(
            vec![Token{ lexeme: "cribe".to_string(), category: Category::Text }],
            buffer.tokens()
        );
    }
}
<|endoftext|>


<commit_before>use std::ops::Deref;
use std::rc::Rc;
use std::cell::RefCell;
use super::Position;
use super::GapBuffer;

/// Read-only wrapper for a `Position`, to allow field level access to a
/// buffer's cursor while simultaneously enforcing bounds-checking when
/// updating its value.
pub struct Cursor {
    pub data: Rc<RefCell<GapBuffer>>,
    pub position: Position,
}

impl Deref for Cursor {
    type Target = Position;

    fn deref(&self) -> &Position {
        &self.position
    }
}

impl Cursor {
    /// Moves the cursor to the specified location. The location is
    /// bounds-checked against the data and the cursor will not be
    /// updated if it is out-of-bounds.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// let in_bounds = scribe::buffer::Position{ line: 0, offset: 2 };
    /// let out_of_bounds = scribe::buffer::Position{ line: 2, offset: 2 };
    /// buffer.insert("scribe");
    ///
    /// buffer.cursor.move_to(in_bounds);
    /// assert_eq!(buffer.cursor.line, 0);
    /// assert_eq!(buffer.cursor.offset, 2);
    ///
    /// buffer.cursor.move_to(out_of_bounds);
    /// assert_eq!(buffer.cursor.line, 0);
    /// assert_eq!(buffer.cursor.offset, 2);
    /// ```
    pub fn move_to(&mut self, position: Position) -> bool {
        if self.data.borrow().in_bounds(&position) {
            self.position = position;
            return true
        }
        false
    }

    /// Decrements the cursor line. The location is bounds-checked against
    /// the data and the cursor will not be updated if it is out-of-bounds.
    pub fn move_up(&mut self) {
        let target_line = self.line-1;
        let mut new_position = Position{ line: target_line, offset: self.offset };

        // Try moving to the same offset on the line above, falling back to its EOL.
        if self.move_to(new_position) == false {
            let mut target_offset = 0;
            for (line_number, line) in self.data.borrow().to_string().lines().enumerate() {
                if line_number == target_line {
                    target_offset = line.len();
                }
            }
            self.move_to(Position{ line: target_line, offset: target_offset });
        }
    }

    /// Increments the cursor line. The location is bounds-checked against
    /// the data and the cursor will not be updated if it is out-of-bounds.
    pub fn move_down(&mut self) {
        let new_position = Position{ line: self.line+1, offset: self.offset };
        self.move_to(new_position);
    }

    /// Decrements the cursor offset. The location is bounds-checked against
    /// the data and the cursor will not be updated if it is out-of-bounds.
    pub fn move_left(&mut self) {
        let new_position = Position{ line: self.line, offset: self.offset-1 };
        self.move_to(new_position);
    }

    /// Increments the cursor offset. The location is bounds-checked against
    /// the data and the cursor will not be updated if it is out-of-bounds.
    pub fn move_right(&mut self) {
        let new_position = Position{ line: self.line, offset: self.offset+1 };
        self.move_to(new_position);
    }
}

#[cfg(test)]
mod tests {
    use super::Cursor;
    use super::super::gap_buffer;
    use super::super::Position;
    use std::rc::Rc;
    use std::cell::RefCell;

    #[test]
    fn move_up_goes_to_EOL_if_offset_would_be_out_of_range() {
        let mut buffer = Rc::new(RefCell::new(gap_buffer::new("This is a test.\nAnother line that is longer.".to_string())));
        let position = Position{ line: 1, offset: 20 };
        let mut cursor = Cursor{ data: buffer, position: position };
        cursor.move_up();
        assert_eq!(cursor.line, 0);
        assert_eq!(cursor.offset, 15);
    }
}
<commit_msg>Remove unused position mutability.<commit_after>use std::ops::Deref;
use std::rc::Rc;
use std::cell::RefCell;
use super::Position;
use super::GapBuffer;

/// Read-only wrapper for a `Position`, to allow field level access to a
/// buffer's cursor while simultaneously enforcing bounds-checking when
/// updating its value.
pub struct Cursor {
    pub data: Rc<RefCell<GapBuffer>>,
    pub position: Position,
}

impl Deref for Cursor {
    type Target = Position;

    fn deref(&self) -> &Position {
        &self.position
    }
}

impl Cursor {
    /// Moves the cursor to the specified location. The location is
    /// bounds-checked against the data and the cursor will not be
    /// updated if it is out-of-bounds.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// let in_bounds = scribe::buffer::Position{ line: 0, offset: 2 };
    /// let out_of_bounds = scribe::buffer::Position{ line: 2, offset: 2 };
    /// buffer.insert("scribe");
    ///
    /// buffer.cursor.move_to(in_bounds);
    /// assert_eq!(buffer.cursor.line, 0);
    /// assert_eq!(buffer.cursor.offset, 2);
    ///
    /// buffer.cursor.move_to(out_of_bounds);
    /// assert_eq!(buffer.cursor.line, 0);
    /// assert_eq!(buffer.cursor.offset, 2);
    /// ```
    pub fn move_to(&mut self, position: Position) -> bool {
        if self.data.borrow().in_bounds(&position) {
            self.position = position;
            return true
        }
        false
    }

    /// Decrements the cursor line. The location is bounds-checked against
    /// the data and the cursor will not be updated if it is out-of-bounds.
    pub fn move_up(&mut self) {
        let target_line = self.line-1;
        let new_position = Position{ line: target_line, offset: self.offset };

        // Try moving to the same offset on the line above, falling back to its EOL.
        if self.move_to(new_position) == false {
            let mut target_offset = 0;
            for (line_number, line) in self.data.borrow().to_string().lines().enumerate() {
                if line_number == target_line {
                    target_offset = line.len();
                }
            }
            self.move_to(Position{ line: target_line, offset: target_offset });
        }
    }

    /// Increments the cursor line. The location is bounds-checked against
    /// the data and the cursor will not be updated if it is out-of-bounds.
    pub fn move_down(&mut self) {
        let new_position = Position{ line: self.line+1, offset: self.offset };
        self.move_to(new_position);
    }

    /// Decrements the cursor offset. The location is bounds-checked against
    /// the data and the cursor will not be updated if it is out-of-bounds.
    pub fn move_left(&mut self) {
        let new_position = Position{ line: self.line, offset: self.offset-1 };
        self.move_to(new_position);
    }

    /// Increments the cursor offset. The location is bounds-checked against
    /// the data and the cursor will not be updated if it is out-of-bounds.
    pub fn move_right(&mut self) {
        let new_position = Position{ line: self.line, offset: self.offset+1 };
        self.move_to(new_position);
    }
}

#[cfg(test)]
mod tests {
    use super::Cursor;
    use super::super::gap_buffer;
    use super::super::Position;
    use std::rc::Rc;
    use std::cell::RefCell;

    #[test]
    fn move_up_goes_to_EOL_if_offset_would_be_out_of_range() {
        let mut buffer = Rc::new(RefCell::new(gap_buffer::new("This is a test.\nAnother line that is longer.".to_string())));
        let position = Position{ line: 1, offset: 20 };
        let mut cursor = Cursor{ data: buffer, position: position };
        cursor.move_up();
        assert_eq!(cursor.line, 0);
        assert_eq!(cursor.offset, 15);
    }
}
<|endoftext|>


<commit_before>use entities::*;
use std::cmp::Ordering;
use super::geo::{self, Coordinate};

trait DistanceTo {
    fn distance_to(&self, &Coordinate) -> f64;
}

impl DistanceTo for Entry {
    fn distance_to(&self, c: &Coordinate) -> f64 {
        geo::distance(&Coordinate {
                          lat: self.lat,
                          lng: self.lng,
                      },
                      c)
    }
}

pub trait SortByDistanceTo {
    fn sort_by_distance_to(&mut self, &Coordinate);
}

impl SortByDistanceTo for Vec<Entry> {
    fn sort_by_distance_to(&mut self, c: &Coordinate) {
        self.sort_by(|a, _| if a.lat.is_finite() && a.lng.is_finite() {
            Ordering::Less
        } else {
            warn!("inivalid coordinate: {}/{}", a.lat, a.lng);
            Ordering::Greater
        });
        self.sort_by(|a, b| {
            a.distance_to(c).partial_cmp(&b.distance_to(c)).unwrap_or(Ordering::Equal)
        })
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    fn new_entry(id: &str, lat: f64, lng: f64) -> Entry { Entry{
            id          : id.into(),
            created     : 0,
            version     : 0,
            title       : "foo".into(),
            description : "bar".into(),
            lat         : lat,
            lng         : lng,
            street      : None,
            zip         : None,
            city        : None,
            country     : None,
            email       : None,
            telephone   : None,
            homepage    : None,
            categories  : vec![],
            license     : None,
        }
    }

    #[test]
    fn sort_by_distance() {
        let mut entries = vec![new_entry("a", 1.0, 0.0),
                               new_entry("b", 0.0, 0.0),
                               new_entry("c", 1.0, 1.0),
                               new_entry("d", 0.0, 0.5),
                               new_entry("e", -1.0, -1.0)];
        let x = Coordinate {
            lat: 0.0,
            lng: 0.0,
        };
        entries.sort_by_distance_to(&x);
        assert_eq!(entries[0].id, "b");
        assert_eq!(entries[1].id, "d");
        assert_eq!(entries[2].id, "a");
        assert_eq!(entries[3].id, "c");
        assert_eq!(entries[4].id, "e");
    }

    use std::f64::{NAN, INFINITY};

    #[test]
    fn sort_with_invalid_coordinates() {
        let mut entries = vec![new_entry("a", 1.0, NAN),
                               new_entry("b", 1.0, INFINITY),
                               new_entry("c", 2.0, 0.0),
                               new_entry("d", NAN, NAN),
                               new_entry("e", 1.0, 0.0)];
        let x = Coordinate {
            lat: 0.0,
            lng: 0.0,
        };
        entries.sort_by_distance_to(&x);
        assert_eq!(entries[0].id, "e");
        assert_eq!(entries[1].id, "c");

        let mut entries =
            vec![new_entry("a", 2.0, 0.0), new_entry("b", 0.0, 0.0), new_entry("c", 1.0, 0.0)];

        let x = Coordinate {
            lat: NAN,
            lng: 0.0,
        };
        entries.sort_by_distance_to(&x);
        assert_eq!(entries[0].id, "a");
        assert_eq!(entries[1].id, "b");
        assert_eq!(entries[2].id, "c");
    }

}
<commit_msg>fix sorting with NaN coordinates<commit_after>use entities::*;
use std::cmp::Ordering;
use super::geo::{self, Coordinate};

trait DistanceTo {
    fn distance_to(&self, &Coordinate) -> f64;
}

impl DistanceTo for Entry {
    fn distance_to(&self, c: &Coordinate) -> f64 {
        geo::distance(&Coordinate {
                          lat: self.lat,
                          lng: self.lng,
                      },
                      c)
    }
}

pub trait SortByDistanceTo {
    fn sort_by_distance_to(&mut self, &Coordinate);
}

impl SortByDistanceTo for Vec<Entry> {
    fn sort_by_distance_to(&mut self, c: &Coordinate) {
        if !(c.lat.is_finite() && c.lng.is_finite()) {
            return;
        }
        self.sort_by(|a, _|
            if a.lat.is_finite() && a.lng.is_finite() {
                Ordering::Less
            } else {
                warn!("inivalid coordinate: {}/{}", a.lat, a.lng);
                Ordering::Greater
            }
        );
        self.sort_by(|a, b| {
            a.distance_to(c).partial_cmp(&b.distance_to(c)).unwrap_or(Ordering::Equal)
        })
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    fn new_entry(id: &str, lat: f64, lng: f64) -> Entry { Entry{
            id          : id.into(),
            created     : 0,
            version     : 0,
            title       : "foo".into(),
            description : "bar".into(),
            lat         : lat,
            lng         : lng,
            street      : None,
            zip         : None,
            city        : None,
            country     : None,
            email       : None,
            telephone   : None,
            homepage    : None,
            categories  : vec![],
            license     : None,
        }
    }

    #[test]
    fn sort_by_distance() {
        let mut entries = vec![new_entry("a", 1.0, 0.0),
                               new_entry("b", 0.0, 0.0),
                               new_entry("c", 1.0, 1.0),
                               new_entry("d", 0.0, 0.5),
                               new_entry("e", -1.0, -1.0)];
        let x = Coordinate {
            lat: 0.0,
            lng: 0.0,
        };
        entries.sort_by_distance_to(&x);
        assert_eq!(entries[0].id, "b");
        assert_eq!(entries[1].id, "d");
        assert_eq!(entries[2].id, "a");
        assert!(entries[3].id == "c" || entries[3].id == "e");
        assert!(entries[4].id == "c" || entries[4].id == "e");
    }

    use std::f64::{NAN, INFINITY};

    #[test]
    fn sort_with_invalid_coordinates() {
        let mut entries = vec![new_entry("a", 1.0, NAN),
                               new_entry("b", 1.0, INFINITY),
                               new_entry("c", 2.0, 0.0),
                               new_entry("d", NAN, NAN),
                               new_entry("e", 1.0, 0.0)];
        let x = Coordinate {
            lat: 0.0,
            lng: 0.0,
        };
        entries.sort_by_distance_to(&x);
        assert_eq!(entries[0].id, "e");
        assert_eq!(entries[1].id, "c");

        let mut entries =
            vec![new_entry("a", 2.0, 0.0), new_entry("b", 0.0, 0.0), new_entry("c", 1.0, 0.0)];

        let x = Coordinate {
            lat: NAN,
            lng: 0.0,
        };
        entries.sort_by_distance_to(&x);
        assert_eq!(entries[0].id, "a");
        assert_eq!(entries[1].id, "b");
        assert_eq!(entries[2].id, "c");
    }

}
<|endoftext|>