<commit_before>#![cfg_attr(feature = "nightly", feature(catch_panic))]
extern crate disque;

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::thread::{spawn, JoinHandle};
#[cfg(feature = "nightly")] use std::thread::catch_panic;

use disque::Disque;

#[derive(Clone)]
pub enum JobStatus {
    FastAck,
    AckJob,
    NAck,
}

pub trait Handler {
    fn process_job(&self, queue_name: &[u8], jobid: &String, body: Vec<u8>) -> JobStatus;
    fn process_error(&self, queue_name: &[u8], jobid: &String, nack: u32, additional_deliveries: u32) -> bool;
}

#[derive(Clone)]
struct HandlerWrapper<H: Handler> {
    handler: Arc<H>,
}
unsafe impl<H: Handler> Send for HandlerWrapper<H> {}
unsafe impl<H: Handler> Sync for HandlerWrapper<H> {}

#[cfg(feature = "nightly")]
macro_rules! spawn {
    ($func: expr, $err: expr) => {
        spawn(move || {
            match catch_panic(move || $func) {
                Ok(_) => (),
                Err(e) => ($err)(e),
            }
        })
    }
}

#[cfg(not(feature = "nightly"))]
macro_rules! spawn {
    ($func: expr, $err: expr) => {
        spawn(move || $func)
    }
}

#[allow(unused_variables)]
fn create_worker<H: Handler + Clone + 'static>(position: usize,
        task_rx: Receiver<Option<(Vec<u8>, String, Vec<u8>, u32, u32)>>,
        completion_tx: Sender<Result<(usize, String, JobStatus), usize>>,
        handler_: HandlerWrapper<H>,
        ) -> JoinHandle<()> {
    let handlerw = handler_.clone();
    let completion_tx2 = completion_tx.clone();
    spawn!({
        let handler = handlerw.handler;
        loop {
            let (queue, jobid, job, nack,
                additional_deliveries) = match task_rx.recv() {
                Ok(o) => match o {
                    Some(v) => v,
                    None => break,
                },
                Err(e) => {
                    // TODO: better log
                    println!("Error in worker thread {:?}", e);
                    break;
                }
            };

            if nack > 0 || additional_deliveries > 0 {
                if !handler.process_error(&*queue, &jobid, nack,
                    additional_deliveries) {
                    return;
                }
            }
            let status = handler.process_job(&*queue, &jobid, job);

            completion_tx.send(Ok((position, jobid, status))).unwrap();
        }
    }, |e| {
        println!("handle panic {:?}", e);
        completion_tx2.send(Err(position)).unwrap();
    })
}

pub struct EventLoop<H: Handler + Clone + 'static> {
    disque: Disque,
    workers: Vec<(JoinHandle<()>, Sender<Option<(Vec<u8>, String, Vec<u8>, u32, u32)>>)>,
    completion_rx: Receiver<Result<(usize, String, JobStatus), usize>>,
    free_workers: HashSet<usize>,
    queues: HashSet<Vec<u8>>,
    hello: (u8, String, Vec<(String, String, u16, u32)>),
    node_counter: HashMap<Vec<u8>, usize>,
    completion_tx: Sender<Result<(usize, String, JobStatus), usize>>,
    ahandler: HandlerWrapper<H>
}

impl<H: Handler + Clone + 'static> EventLoop<H> {
    pub fn new(
            disque: Disque, numworkers: usize,
            handler: H) -> Self {
        let mut workers = Vec::with_capacity(numworkers);
        let mut free_workers = HashSet::with_capacity(numworkers);
        let (completion_tx, completion_rx) = channel();
        let ahandler = HandlerWrapper { handler: Arc::new(handler) };
        for i in 0..numworkers {
            let (task_tx, task_rx) = channel();
            let jg = create_worker(i, task_rx, completion_tx.clone(),
                    ahandler.clone());
            workers.push((jg, task_tx));
            free_workers.insert(i);
        }
        let hello = disque.hello().unwrap();
        EventLoop {
            disque: disque,
            completion_rx: completion_rx,
            workers: workers,
            hello: hello,
            free_workers: free_workers,
            queues: HashSet::new(),
            node_counter: HashMap::new(),
            completion_tx: completion_tx,
            ahandler: ahandler,
        }
    }

    pub fn watch_queue(&mut self, queue_name: Vec<u8>) {
        self.queues.insert(queue_name);
    }

    pub fn unwatch_queue(&mut self, queue_name: &Vec<u8>) {
        self.queues.remove(queue_name);
    }

    fn completed(&mut self, worker: usize, jobid: String, status: JobStatus) {
        self.free_workers.insert(worker);
        match status {
            JobStatus::FastAck => self.disque.fastackjob(jobid.as_bytes()),
            JobStatus::AckJob => self.disque.ackjob(jobid.as_bytes()),
            JobStatus::NAck => self.disque.nackjob(jobid.as_bytes()),
        }.unwrap();
    }

    fn handle_worker_panic(&mut self, worker: usize) {
        if self.workers.len() == 0 {
            // shutting down
            return;
        }

        let (task_tx, task_rx) = channel();
        let jg = create_worker(worker, task_rx, self.completion_tx.clone(),
                self.ahandler.clone());
        self.workers[worker] = (jg, task_tx);
        self.free_workers.insert(worker);
    }

    fn mark_completed(&mut self, blocking: bool) {
        macro_rules! recv {
            ($func: ident) => {
                match self.completion_rx.$func() {
                    Ok(c) => match c {
                        Ok(c) => self.completed(c.0, c.1, c.2),
                        Err(worker) => self.handle_worker_panic(worker),
                    },
                    Err(_) => return,
                }
            }
        }

        if blocking {
            recv!(recv);
        }
        loop {
            recv!(try_recv);
        }
    }

    pub fn choose_favorite_node(&self) -> (Vec<u8>, usize) {
        let mut r = (&Vec::new(), &0);
        for n in self.node_counter.iter() {
            if n.1 >= r.1 {
                r = n;
            }
        }
        (r.0.clone(), r.1.clone())
    }

    pub fn jobcount_current_node(&self) -> usize {
        let nodeid = self.hello.1.as_bytes()[0..8].to_vec();
        self.node_counter.get(&nodeid).unwrap_or(&0).clone()
    }

    pub fn current_node_id(&self) -> String {
        self.hello.1.clone()
    }

    fn run_once(&mut self) -> bool {
        self.mark_completed(false);
        let worker = match self.free_workers.iter().next() {
            Some(w) => w.clone(),
            None => return false,
        };

        let job = match self.disque.getjob_withcounters(false, None,
                &*self.queues.iter().map(|k| &**k).collect::<Vec<_>>()
                ).unwrap() {
            Some(j) => j,
            None => return false,
        };

        let nodeid = job.1.as_bytes()[2..10].to_vec();
        let v = self.node_counter.remove(&nodeid).unwrap_or(0);
        self.node_counter.insert(nodeid, v + 1);

        self.free_workers.remove(&worker);
        self.workers[worker].1.send(Some(job)).unwrap();
        true
    }

    fn connect_to_node(&mut self, new_master: Vec<u8>) -> bool {
        let mut hello = None;
        for node in self.hello.2.iter() {
            if node.0.as_bytes()[..new_master.len()] == *new_master {
                match Disque::open(&*format!("redis://{}:{}/", node.1, node.2)) {
                    Ok(disque) => {
                        hello = Some(match disque.hello() {
                            Ok(hello) => hello,
                            Err(_) => break,
                        });
                        self.disque = disque;
                        break;
                    },
                    Err(_) => (),
                }
                break;
            }
        }
        match hello {
            Some(h) => { self.hello = h; true }
            None => false,
        }
    }

    pub fn do_cycle(&mut self) {
        let (fav_node, fav_count) = self.choose_favorite_node();
        let current_count = self.jobcount_current_node();
        // only change if it is at least 20% better than the current node
        if fav_count as f64 / current_count as f64 > 1.2 {
            self.connect_to_node(fav_node);
        }
    }

    pub fn run(&mut self, cycle: usize) {
        self.run_times_cycle(0, cycle)
    }

    pub fn run_times(&mut self, times: usize) {
        self.run_times_cycle(times, 0)
    }

    pub fn run_times_cycle(&mut self, times: usize, cycle: usize) {
        let mut c = 0;
        let mut counter = 0;
        loop {
            let did_run = self.run_once();
            if did_run {
                if times > 0 {
                    counter += 1;
                    if counter == times {
                        break;
                    }
                }
                if cycle > 0 {
                    c += 1;
                    if c == cycle {
                        self.do_cycle();
                        c = 0;
                    }
                }
            } else {
                self.mark_completed(true);
            }
        }
        self.mark_completed(false);
    }

    pub fn stop(mut self) {
        for worker in std::mem::replace(&mut self.workers, vec![]).into_iter() {
            worker.1.send(None).unwrap();
            worker.0.join().unwrap();
        }
        self.mark_completed(false);
    }
}

#[test]
fn favorite() {
    #[derive(Clone)]
    struct MyHandler;
    impl Handler for MyHandler {
        fn process_job(&self, _: &[u8], _: &String, _: Vec<u8>) -> JobStatus {
            JobStatus::AckJob
        }
        fn process_error(&self, _: &[u8], _: &String, _: u32, _: u32) -> bool {
            false
        }
    }
    let disque = Disque::open("redis://127.0.0.1:7711/").unwrap();
    let mut el = EventLoop::new(disque, 1, MyHandler);
    el.node_counter.insert(vec![1, 2, 3], 123);
    el.node_counter.insert(vec![4, 5, 6], 456);
    el.node_counter.insert(vec![0, 0, 0], 0);
    assert_eq!(el.choose_favorite_node(), (vec![4, 5, 6], 456));
}
<commit_msg>Use JobUpdate for worker result<commit_after>#![cfg_attr(feature = "nightly", feature(catch_panic))]
extern crate disque;

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::thread::{spawn, JoinHandle};
#[cfg(feature = "nightly")] use std::thread::catch_panic;

use disque::Disque;

#[derive(Clone)]
pub enum JobStatus {
    FastAck,
    AckJob,
    NAck,
}

pub trait Handler {
    fn process_job(&self, queue_name: &[u8], jobid: &String, body: Vec<u8>) -> JobStatus;
    fn process_error(&self, queue_name: &[u8], jobid: &String, nack: u32, additional_deliveries: u32) -> bool;
}

#[derive(Clone)]
struct HandlerWrapper<H: Handler> {
    handler: Arc<H>,
}
unsafe impl<H: Handler> Send for HandlerWrapper<H> {}
unsafe impl<H: Handler> Sync for HandlerWrapper<H> {}

#[cfg(feature = "nightly")]
macro_rules! spawn {
    ($func: expr, $err: expr) => {
        spawn(move || {
            match catch_panic(move || $func) {
                Ok(_) => (),
                Err(e) => ($err)(e),
            }
        })
    }
}

#[cfg(not(feature = "nightly"))]
macro_rules! spawn {
    ($func: expr, $err: expr) => {
        spawn(move || $func)
    }
}

#[allow(dead_code)]
enum JobUpdate {
    Success(usize, String, JobStatus),
    Failure(usize),
}

#[allow(unused_variables)]
fn create_worker<H: Handler + Clone + 'static>(position: usize,
        task_rx: Receiver<Option<(Vec<u8>, String, Vec<u8>, u32, u32)>>,
        completion_tx: Sender<JobUpdate>,
        handler_: HandlerWrapper<H>,
        ) -> JoinHandle<()> {
    let handlerw = handler_.clone();
    let completion_tx2 = completion_tx.clone();
    spawn!({
        let handler = handlerw.handler;
        loop {
            let (queue, jobid, job, nack,
                additional_deliveries) = match task_rx.recv() {
                Ok(o) => match o {
                    Some(v) => v,
                    None => break,
                },
                Err(e) => {
                    // TODO: better log
                    println!("Error in worker thread {:?}", e);
                    break;
                }
            };

            if nack > 0 || additional_deliveries > 0 {
                if !handler.process_error(&*queue, &jobid, nack,
                    additional_deliveries) {
                    return;
                }
            }
            let status = handler.process_job(&*queue, &jobid, job);

            completion_tx.send(JobUpdate::Success(position, jobid, status)).unwrap();
        }
    }, |e| {
        println!("handle panic {:?}", e);
        completion_tx2.send(JobUpdate::Failure(position)).unwrap();
    })
}

pub struct EventLoop<H: Handler + Clone + 'static> {
    disque: Disque,
    workers: Vec<(JoinHandle<()>, Sender<Option<(Vec<u8>, String, Vec<u8>, u32, u32)>>)>,
    completion_rx: Receiver<JobUpdate>,
    free_workers: HashSet<usize>,
    queues: HashSet<Vec<u8>>,
    hello: (u8, String, Vec<(String, String, u16, u32)>),
    node_counter: HashMap<Vec<u8>, usize>,
    completion_tx: Sender<JobUpdate>,
    ahandler: HandlerWrapper<H>
}

impl<H: Handler + Clone + 'static> EventLoop<H> {
    pub fn new(
            disque: Disque, numworkers: usize,
            handler: H) -> Self {
        let mut workers = Vec::with_capacity(numworkers);
        let mut free_workers = HashSet::with_capacity(numworkers);
        let (completion_tx, completion_rx) = channel();
        let ahandler = HandlerWrapper { handler: Arc::new(handler) };
        for i in 0..numworkers {
            let (task_tx, task_rx) = channel();
            let jg = create_worker(i, task_rx, completion_tx.clone(),
                    ahandler.clone());
            workers.push((jg, task_tx));
            free_workers.insert(i);
        }
        let hello = disque.hello().unwrap();
        EventLoop {
            disque: disque,
            completion_rx: completion_rx,
            workers: workers,
            hello: hello,
            free_workers: free_workers,
            queues: HashSet::new(),
            node_counter: HashMap::new(),
            completion_tx: completion_tx,
            ahandler: ahandler,
        }
    }

    pub fn watch_queue(&mut self, queue_name: Vec<u8>) {
        self.queues.insert(queue_name);
    }

    pub fn unwatch_queue(&mut self, queue_name: &Vec<u8>) {
        self.queues.remove(queue_name);
    }

    fn completed(&mut self, worker: usize, jobid: String, status: JobStatus) {
        self.free_workers.insert(worker);
        match status {
            JobStatus::FastAck => self.disque.fastackjob(jobid.as_bytes()),
            JobStatus::AckJob => self.disque.ackjob(jobid.as_bytes()),
            JobStatus::NAck => self.disque.nackjob(jobid.as_bytes()),
        }.unwrap();
    }

    fn handle_worker_panic(&mut self, worker: usize) {
        if self.workers.len() == 0 {
            // shutting down
            return;
        }

        let (task_tx, task_rx) = channel();
        let jg = create_worker(worker, task_rx, self.completion_tx.clone(),
                self.ahandler.clone());
        self.workers[worker] = (jg, task_tx);
        self.free_workers.insert(worker);
    }

    fn mark_completed(&mut self, blocking: bool) {
        macro_rules! recv {
            ($func: ident) => {
                match self.completion_rx.$func() {
                    Ok(c) => match c {
                        JobUpdate::Success(worker, jobid, status) => {
                            self.completed(worker, jobid, status);
                        },
                        JobUpdate::Failure(worker) => self.handle_worker_panic(worker),
                    },
                    Err(_) => return,
                }
            }
        }

        if blocking {
            recv!(recv);
        }
        loop {
            recv!(try_recv);
        }
    }

    pub fn choose_favorite_node(&self) -> (Vec<u8>, usize) {
        let mut r = (&Vec::new(), &0);
        for n in self.node_counter.iter() {
            if n.1 >= r.1 {
                r = n;
            }
        }
        (r.0.clone(), r.1.clone())
    }

    pub fn jobcount_current_node(&self) -> usize {
        let nodeid = self.hello.1.as_bytes()[0..8].to_vec();
        self.node_counter.get(&nodeid).unwrap_or(&0).clone()
    }

    pub fn current_node_id(&self) -> String {
        self.hello.1.clone()
    }

    fn run_once(&mut self) -> bool {
        self.mark_completed(false);
        let worker = match self.free_workers.iter().next() {
            Some(w) => w.clone(),
            None => return false,
        };

        let job = match self.disque.getjob_withcounters(false, None,
                &*self.queues.iter().map(|k| &**k).collect::<Vec<_>>()
                ).unwrap() {
            Some(j) => j,
            None => return false,
        };

        let nodeid = job.1.as_bytes()[2..10].to_vec();
        let v = self.node_counter.remove(&nodeid).unwrap_or(0);
        self.node_counter.insert(nodeid, v + 1);

        self.free_workers.remove(&worker);
        self.workers[worker].1.send(Some(job)).unwrap();
        true
    }

    fn connect_to_node(&mut self, new_master: Vec<u8>) -> bool {
        let mut hello = None;
        for node in self.hello.2.iter() {
            if node.0.as_bytes()[..new_master.len()] == *new_master {
                match Disque::open(&*format!("redis://{}:{}/", node.1, node.2)) {
                    Ok(disque) => {
                        hello = Some(match disque.hello() {
                            Ok(hello) => hello,
                            Err(_) => break,
                        });
                        self.disque = disque;
                        break;
                    },
                    Err(_) => (),
                }
                break;
            }
        }
        match hello {
            Some(h) => { self.hello = h; true }
            None => false,
        }
    }

    pub fn do_cycle(&mut self) {
        let (fav_node, fav_count) = self.choose_favorite_node();
        let current_count = self.jobcount_current_node();
        // only change if it is at least 20% better than the current node
        if fav_count as f64 / current_count as f64 > 1.2 {
            self.connect_to_node(fav_node);
        }
    }

    pub fn run(&mut self, cycle: usize) {
        self.run_times_cycle(0, cycle)
    }

    pub fn run_times(&mut self, times: usize) {
        self.run_times_cycle(times, 0)
    }

    pub fn run_times_cycle(&mut self, times: usize, cycle: usize) {
        let mut c = 0;
        let mut counter = 0;
        loop {
            let did_run = self.run_once();
            if did_run {
                if times > 0 {
                    counter += 1;
                    if counter == times {
                        break;
                    }
                }
                if cycle > 0 {
                    c += 1;
                    if c == cycle {
                        self.do_cycle();
                        c = 0;
                    }
                }
            } else {
                self.mark_completed(true);
            }
        }
        self.mark_completed(false);
    }

    pub fn stop(mut self) {
        for worker in std::mem::replace(&mut self.workers, vec![]).into_iter() {
            worker.1.send(None).unwrap();
            worker.0.join().unwrap();
        }
        self.mark_completed(false);
    }
}

#[test]
fn favorite() {
    #[derive(Clone)]
    struct MyHandler;
    impl Handler for MyHandler {
        fn process_job(&self, _: &[u8], _: &String, _: Vec<u8>) -> JobStatus {
            JobStatus::AckJob
        }
        fn process_error(&self, _: &[u8], _: &String, _: u32, _: u32) -> bool {
            false
        }
    }
    let disque = Disque::open("redis://127.0.0.1:7711/").unwrap();
    let mut el = EventLoop::new(disque, 1, MyHandler);
    el.node_counter.insert(vec![1, 2, 3], 123);
    el.node_counter.insert(vec![4, 5, 6], 456);
    el.node_counter.insert(vec![0, 0, 0], 0);
    assert_eq!(el.choose_favorite_node(), (vec![4, 5, 6], 456));
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::borrow::Borrow;
use std::hash::{Hash,Hasher};

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Param {
    Short(char),
    Long(String),
}

impl Borrow<str> for Param {
    fn borrow(&self) -> &str {
        if let Param::Long(ref string) = *self {
            string
        } else {
            ""
        }
    }
}

impl Borrow<char> for Param {
    fn borrow(&self) -> &char {
        if let Param::Short(ref ch) = *self {
            ch
        } else {
            const CH: &'static char = &'\0';
            CH
        }
    }
}

impl Hash for Param {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match *self {
            Param::Short(ref c) => c.hash(state),
            Param::Long(ref s) => s.hash(state)
        }
    }
}

#[derive(Clone, Debug, Default, Eq, PartialEq, Hash)]
struct Rhs<T> {
    value: T,
    occurrences: usize,
}

impl<T> Rhs<T> {
    fn new<U: Into<T>>(value: U) -> Self {
        Rhs { value: value.into(), occurrences: 0 }
    }
}

#[derive(Clone, Debug, Eq, PartialEq, Hash)]
enum OptRhs {
    With(Rhs<String>, bool),
    Empty,
}

#[derive(Clone, Debug, Eq, PartialEq, Hash)]
enum Value {
    Flag(Rhs<bool>),
    Opt(OptRhs),
}

/// Our homebrewed Arg Parser
#[derive(Clone, Debug, Default)]
pub struct ArgParser {
    params: HashMap<Param, Value>,
    invalid: Vec<Param>,
    pub args: Vec<String>,
}

impl ArgParser {
    /// Create a new ArgParser object
    ///
    /// The capacity specifies the initial capacity of the number parameters.
    /// Always good to set it at the number of flags and opts total.
    pub fn new(capacity: usize) -> Self {
        ArgParser {
            params: HashMap::with_capacity(capacity),
            invalid: Vec::new(),
            args: Vec::new(),
        }
    }

    /// Builder method for adding both short and long flags
    ///
    /// Flags are just parameters that have no assigned values. They are used
    /// for when certain features or options have been enabled for the application
    ///
    /// For example
    /// > ls -l --human-readable
    ///   ^  ^  ^
    ///   |  |  |
    ///   |  |  `-- A long flag to enable human readable numbers.
    ///   |  `-- A short flag to enable the long format.
    ///   `-- The command to list files.
    pub fn add_flag(mut self, short: &str, long: &str) -> Self {
        if let Some(short) = short.chars().next() {
            self.params.insert(Param::Short(short), Value::Flag(Rhs::default()));
        }
        if !long.is_empty() {
            self.params.insert(Param::Long(long.to_owned()), Value::Flag(Rhs::default()));
        }
        self
    }

    /// Builder method for adding both short and long opts
    ///
    /// Opts are parameters that hold assigned values. They are used
    /// for when certain features or options have been enabled for the application
    ///
    /// For example
    /// > ls -T 4 --color=always
    ///   ^  ^    ^
    ///   |  |    |
    ///   |  |    `-- A long opt to enable the use of color with value `always`.
    ///   |  `-- A short opt to set tab size to the value `4`.
    ///   `-- The command to list files.
    pub fn add_opt(mut self, short: &str, long: &str) -> Self {
        if let Some(short) = short.chars().next() {
            self.params.insert(Param::Short(short), Value::Opt(OptRhs::Empty));
        }
        if !long.is_empty() {
            self.params.insert(Param::Long(long.to_owned()), Value::Opt(OptRhs::Empty));
        }
        self
    }

    pub fn add_opt_default(mut self, short: &str, long: &str, default: &str) -> Self {
        if let Some(short) = short.chars().next() {
            self.params.insert(Param::Short(short), Value::Opt(OptRhs::With(Rhs::new(default), false)));
        }
        if !long.is_empty() {
            self.params.insert(Param::Long(long.to_owned()), Value::Opt(OptRhs::With(Rhs::new(default), false)));
        }
        self
    }

    /// Start parsing user inputted args for which flags and opts are used at
    /// runtime. The rest of the args that are not associated to opts get added
    /// to `ArgParser.args`.
    pub fn initialize<A: Iterator<Item=String>>(&mut self, args: A) {
        let mut args = args.skip(1);
        while let Some(arg) = args.next() {
            if arg.starts_with("--") {
                // Remove both dashes
                let arg = &arg[2..];
                if arg.is_empty() {
                    //Arg `--` means we are done parsing args, collect the rest
                    self.args.extend(args);
                    break;
                }
                if let Some(i) = arg.find('=') {
                    let (lhs, rhs) = arg.split_at(i);
                    let rhs = &rhs[1..]; // slice off the `=` char
                    match self.params.get_mut(lhs) {
                        Some(&mut Value::Opt(ref mut value)) => {
                            match value {
                                &mut OptRhs::With(ref mut opt_rhs, ref mut occur) => {
                                    opt_rhs.value.clear();
                                    opt_rhs.value.push_str(rhs);
                                    opt_rhs.occurrences += 1;
                                    *occur = true;
                                }
                                &mut OptRhs::Empty => {
                                    let mut rhs = Rhs::new(rhs);
                                    rhs.occurrences = 1;
                                    *value = OptRhs::With(rhs, true);
                                }
                            }
                        }
                        _ => self.invalid.push(Param::Long(lhs.to_owned())),
                    }
                }
                else {
                    match self.params.get_mut(arg) {
                        Some(&mut Value::Flag(ref mut rhs)) => {
                            rhs.value = true;
                            rhs.occurrences += 1;
                        }
                        Some(&mut Value::Opt(OptRhs::With(ref mut rhs, ref mut occur))) => {
                            rhs.occurrences += 1;
                            *occur = true;
                        }
                        _ => self.invalid.push(Param::Long(arg.to_owned())),
                    }
                }
            }
            else if arg.starts_with("-") {
                let mut chars = arg[1..].chars();
                while let Some(ch) = chars.next() {
                    match self.params.get_mut(&ch) {
                        Some(&mut Value::Flag(ref mut rhs)) => {
                            rhs.value = true;
                            rhs.occurrences += 1;
                        }
                        Some(&mut Value::Opt(OptRhs::With(ref mut rhs, ref mut occur))) => {
                            let rest: String = chars.collect();
                            if !rest.is_empty() {
                                rhs.value = rest;
                                *occur = true;
                            } else if let Some(arg) = args.next() {
                                rhs.value = arg;
                                *occur = true;
                            }
                            break;
                        },
                        _ => self.invalid.push(Param::Short(ch)),
                    }
                }
            }
            else {
                self.args.push(arg);
            }
        }
    }

    /// Check if a Flag or Opt has been found after initialization.
    pub fn flagged
(&self, name: &P) -> bool
        where Param: Borrow

    {
        match self.params.get(name) {
            Some(&Value::Flag(ref rhs)) => rhs.value,
            Some(&Value::Opt(OptRhs::With(_, occur))) => occur,
            _ => false,
        }
    }

    /// Modify the state of a flag. Use `true` if the flag is to be enabled. Use `false` to
    /// disable its use.
    pub fn set_flag<F: Hash + Eq + ?Sized>(&mut self, flag: &F, state: bool)
        where Param: Borrow<F>
    {
        if let Some(&mut Value::Flag(ref mut rhs)) = self.params.get_mut(flag) {
            rhs.value = state;
        }
    }

    /// Modify the state value of an opt. Use `Some(String)` to set if the opt is to be enabled and
    /// has been assigned a value from `String`. Use `None` to disable the opt's use.
    pub fn set_opt<O: Hash + Eq + ?Sized>(&mut self, opt: &O, state: Option<String>)
        where Param: Borrow<O>
    {
        if let Some(&mut Value::Opt(OptRhs::With(ref mut rhs, ref mut occur))) = self.params.get_mut(opt) {
            match state {
                Some(input) => {
                    rhs.value = input;
                    *occur = true;
                }
                None => *occur = false,
            }
        }
    }

    /// Get the value of an Opt. If it has been set or defaulted, it will return a `Some(String)`
    /// value otherwise it will return None.
    pub fn get_opt<O: Hash + Eq + ?Sized>(&self, opt: &O) -> Option<String>
        where Param: Borrow<O>
    {
        if let Some(&Value::Opt(OptRhs::With(ref rhs, _))) = self.params.get(opt) {
            return Some(rhs.value.clone());
        }
        None
    }

    pub fn flagged_invalid(&self) -> Result<(), String> {
        if self.invalid.is_empty() {
            return Ok(());
        }

        let mut and: bool = false;
        let mut output =
            if self.invalid.len() == 1 {
                "Invalid parameter"
            } else {
                and = true;
                "Invalid parameters"
            }.to_owned();

        let mut iter = self.invalid.iter().peekable();
        while let Some(param) = iter.next() {
            match param {
                &Param::Short(ch) => {
                    output += " '-";
                    output.push(ch);
                    output.push('\'');
                }
                &Param::Long(ref s) => {
                    output += " '--";
                    output += s;
                    output.push('\'');
                }
            }
            if and && iter.peek().is_some() {
                output += " and";
            }
        }
        output.push('\n');
        Err(output)
    }
}

pub fn to_human_readable_string(size: u64) -> String {
    if size < 1024 {
        return format!("{}", size);
    }

    static UNITS: [&'static str; 7] = ["", "K", "M", "G", "T", "P", "E"];

    let sizef = size as f64;
    let digit_groups = (sizef.log10() / 1024f64.log10()) as i32;
    format!("{:.1}{}",
            sizef / 1024f64.powf(digit_groups as f64),
            UNITS[digit_groups as usize])
}

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

    #[test]
    fn stop_parsing() {
        let args = vec![String::from("binname"), String::from("-a"), String::from("--"), String::from("-v")];
        let mut parser = ArgParser::new(2);
        parser = parser.add_flag("a", "")
                       .add_flag("v", "");
        parser.initialize(args.into_iter());
        assert!(parser.flagged(&'a'));
        assert!(!parser.flagged(&'v'));
        assert!(parser.args[0] == "-v");
    }

    #[test]
    fn short_opts() {
        let args = vec![String::from("binname"), String::from("-asdf"), String::from("-f"), String::from("foo")];
        let mut parser = ArgParser::new(4);
        parser = parser.add_flag("a", "")
                       .add_flag("d", "")
                       .add_opt("s", "")
                       .add_opt("f", "");
        parser.initialize(args.into_iter());
        assert!(parser.flagged(&'a'));
        assert!(!parser.flagged(&'d'));
        assert!(parser.get_opt(&'s') == Some(String::from("df")));
        assert!(parser.get_opt(&'f') == Some(String::from("foo")));
    }

    #[test]
    fn long_opts() {
        let args = vec![String::from("binname"), String::from("--foo=bar")];
        let mut parser = ArgParser::new(4);
        parser = parser.add_opt("", "foo");
        parser.initialize(args.into_iter());
        assert!(parser.get_opt("foo") == Some(String::from("bar")));
    }
}
<commit_msg>Fix cargo test failure<commit_after>use std::collections::HashMap;
use std::borrow::Borrow;
use std::hash::{Hash,Hasher};

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Param {
    Short(char),
    Long(String),
}

impl Borrow<str> for Param {
    fn borrow(&self) -> &str {
        if let Param::Long(ref string) = *self {
            string
        } else {
            ""
        }
    }
}

impl Borrow<char> for Param {
    fn borrow(&self) -> &char {
        if let Param::Short(ref ch) = *self {
            ch
        } else {
            const CH: &'static char = &'\0';
            CH
        }
    }
}

impl Hash for Param {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match *self {
            Param::Short(ref c) => c.hash(state),
            Param::Long(ref s) => s.hash(state)
        }
    }
}

#[derive(Clone, Debug, Default, Eq, PartialEq, Hash)]
struct Rhs<T> {
    value: T,
    occurrences: usize,
}

impl<T> Rhs<T> {
    fn new<U: Into<T>>(value: U) -> Self {
        Rhs { value: value.into(), occurrences: 0 }
    }
}

#[derive(Clone, Debug, Eq, PartialEq, Hash)]
enum OptRhs {
    With(Rhs<String>, bool),
    Empty,
}

#[derive(Clone, Debug, Eq, PartialEq, Hash)]
enum Value {
    Flag(Rhs<bool>),
    Opt(OptRhs),
}

/// Our homebrewed Arg Parser
#[derive(Clone, Debug, Default)]
pub struct ArgParser {
    params: HashMap<Param, Value>,
    invalid: Vec<Param>,
    pub args: Vec<String>,
}

impl ArgParser {
    /// Create a new ArgParser object
    ///
    /// The capacity specifies the initial capacity of the number parameters.
    /// Always good to set it at the number of flags and opts total.
    pub fn new(capacity: usize) -> Self {
        ArgParser {
            params: HashMap::with_capacity(capacity),
            invalid: Vec::new(),
            args: Vec::new(),
        }
    }

    /// Builder method for adding both short and long flags
    ///
    /// Flags are just parameters that have no assigned values. They are used
    /// for when certain features or options have been enabled for the application
    ///
    /// For example
    /// > ls -l --human-readable
    ///   ^  ^  ^
    ///   |  |  |
    ///   |  |  `-- A long flag to enable human readable numbers.
    ///   |  `-- A short flag to enable the long format.
    ///   `-- The command to list files.
    pub fn add_flag(mut self, short: &str, long: &str) -> Self {
        if let Some(short) = short.chars().next() {
            self.params.insert(Param::Short(short), Value::Flag(Rhs::default()));
        }
        if !long.is_empty() {
            self.params.insert(Param::Long(long.to_owned()), Value::Flag(Rhs::default()));
        }
        self
    }

    /// Builder method for adding both short and long opts
    ///
    /// Opts are parameters that hold assigned values. They are used
    /// for when certain features or options have been enabled for the application
    ///
    /// For example
    /// > ls -T 4 --color=always
    ///   ^  ^    ^
    ///   |  |    |
    ///   |  |    `-- A long opt to enable the use of color with value `always`.
    ///   |  `-- A short opt to set tab size to the value `4`.
    ///   `-- The command to list files.
    pub fn add_opt(mut self, short: &str, long: &str) -> Self {
        if let Some(short) = short.chars().next() {
            self.params.insert(Param::Short(short), Value::Opt(OptRhs::Empty));
        }
        if !long.is_empty() {
            self.params.insert(Param::Long(long.to_owned()), Value::Opt(OptRhs::Empty));
        }
        self
    }

    pub fn add_opt_default(mut self, short: &str, long: &str, default: &str) -> Self {
        if let Some(short) = short.chars().next() {
            self.params.insert(Param::Short(short), Value::Opt(OptRhs::With(Rhs::new(default), false)));
        }
        if !long.is_empty() {
            self.params.insert(Param::Long(long.to_owned()), Value::Opt(OptRhs::With(Rhs::new(default), false)));
        }
        self
    }

    /// Start parsing user inputted args for which flags and opts are used at
    /// runtime. The rest of the args that are not associated to opts get added
    /// to `ArgParser.args`.
    pub fn initialize<A: Iterator<Item=String>>(&mut self, args: A) {
        let mut args = args.skip(1);
        while let Some(arg) = args.next() {
            if arg.starts_with("--") {
                // Remove both dashes
                let arg = &arg[2..];
                if arg.is_empty() {
                    //Arg `--` means we are done parsing args, collect the rest
                    self.args.extend(args);
                    break;
                }
                if let Some(i) = arg.find('=') {
                    let (lhs, rhs) = arg.split_at(i);
                    let rhs = &rhs[1..]; // slice off the `=` char
                    match self.params.get_mut(lhs) {
                        Some(&mut Value::Opt(ref mut value)) => {
                            match value {
                                &mut OptRhs::With(ref mut opt_rhs, ref mut occur) => {
                                    opt_rhs.value.clear();
                                    opt_rhs.value.push_str(rhs);
                                    opt_rhs.occurrences += 1;
                                    *occur = true;
                                }
                                &mut OptRhs::Empty => {
                                    let mut rhs = Rhs::new(rhs);
                                    rhs.occurrences = 1;
                                    *value = OptRhs::With(rhs, true);
                                }
                            }
                        }
                        _ => self.invalid.push(Param::Long(lhs.to_owned())),
                    }
                }
                else {
                    match self.params.get_mut(arg) {
                        Some(&mut Value::Flag(ref mut rhs)) => {
                            rhs.value = true;
                            rhs.occurrences += 1;
                        }
                        Some(&mut Value::Opt(OptRhs::With(ref mut rhs, ref mut occur))) => {
                            rhs.occurrences += 1;
                            *occur = true;
                        }
                        _ => self.invalid.push(Param::Long(arg.to_owned())),
                    }
                }
            }
            else if arg.starts_with("-") {
                let mut chars = arg[1..].chars();
                while let Some(ch) = chars.next() {
                    match self.params.get_mut(&ch) {
                        Some(&mut Value::Flag(ref mut rhs)) => {
                            rhs.value = true;
                            rhs.occurrences += 1;
                        }
                        Some(&mut Value::Opt(ref mut opt_rhs)) => {
                            let rest: String = chars.collect();
                            if !rest.is_empty() {
                                *opt_rhs = OptRhs::With(Rhs::new(rest), true);
                            } else {
                                *opt_rhs = args.next().map(|a| OptRhs::With(Rhs::new(a), true)).unwrap_or(OptRhs::Empty);
                            }
                            break;
                        }
                        None => self.invalid.push(Param::Short(ch)),
                    }
                }
            }
            else {
                self.args.push(arg);
            }
        }
    }

    /// Check if a Flag or Opt has been found after initialization.
    pub fn flagged
(&self, name: &P) -> bool
        where Param: Borrow

    {
        match self.params.get(name) {
            Some(&Value::Flag(ref rhs)) => rhs.value,
            Some(&Value::Opt(OptRhs::With(_, occur))) => occur,
            _ => false,
        }
    }

    /// Modify the state of a flag. Use `true` if the flag is to be enabled. Use `false` to
    /// disable its use.
    pub fn set_flag<F: Hash + Eq + ?Sized>(&mut self, flag: &F, state: bool)
        where Param: Borrow<F>
    {
        if let Some(&mut Value::Flag(ref mut rhs)) = self.params.get_mut(flag) {
            rhs.value = state;
        }
    }

    /// Modify the state value of an opt. Use `Some(String)` to set if the opt is to be enabled and
    /// has been assigned a value from `String`. Use `None` to disable the opt's use.
    pub fn set_opt<O: Hash + Eq + ?Sized>(&mut self, opt: &O, state: Option<String>)
        where Param: Borrow<O>
    {
        if let Some(&mut Value::Opt(OptRhs::With(ref mut rhs, ref mut occur))) = self.params.get_mut(opt) {
            match state {
                Some(input) => {
                    rhs.value = input;
                    *occur = true;
                }
                None => *occur = false,
            }
        }
    }

    /// Get the value of an Opt. If it has been set or defaulted, it will return a `Some(String)`
    /// value otherwise it will return None.
    pub fn get_opt<O: Hash + Eq + ?Sized>(&self, opt: &O) -> Option<String>
        where Param: Borrow<O>
    {
        if let Some(&Value::Opt(OptRhs::With(ref rhs, _))) = self.params.get(opt) {
            return Some(rhs.value.clone());
        }
        None
    }

    pub fn flagged_invalid(&self) -> Result<(), String> {
        if self.invalid.is_empty() {
            return Ok(());
        }

        let mut and: bool = false;
        let mut output =
            if self.invalid.len() == 1 {
                "Invalid parameter"
            } else {
                and = true;
                "Invalid parameters"
            }.to_owned();

        let mut iter = self.invalid.iter().peekable();
        while let Some(param) = iter.next() {
            match param {
                &Param::Short(ch) => {
                    output += " '-";
                    output.push(ch);
                    output.push('\'');
                }
                &Param::Long(ref s) => {
                    output += " '--";
                    output += s;
                    output.push('\'');
                }
            }
            if and && iter.peek().is_some() {
                output += " and";
            }
        }
        output.push('\n');
        Err(output)
    }
}

pub fn to_human_readable_string(size: u64) -> String {
    if size < 1024 {
        return format!("{}", size);
    }

    static UNITS: [&'static str; 7] = ["", "K", "M", "G", "T", "P", "E"];

    let sizef = size as f64;
    let digit_groups = (sizef.log10() / 1024f64.log10()) as i32;
    format!("{:.1}{}",
            sizef / 1024f64.powf(digit_groups as f64),
            UNITS[digit_groups as usize])
}

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

    #[test]
    fn stop_parsing() {
        let args = vec![String::from("binname"), String::from("-a"), String::from("--"), String::from("-v")];
        let mut parser = ArgParser::new(2);
        parser = parser.add_flag("a", "")
                       .add_flag("v", "");
        parser.initialize(args.into_iter());
        assert!(parser.flagged(&'a'));
        assert!(!parser.flagged(&'v'));
        assert!(parser.args[0] == "-v");
    }

    #[test]
    fn short_opts() {
        let args = vec![String::from("binname"), String::from("-asdf"), String::from("-f"), String::from("foo")];
        let mut parser = ArgParser::new(4);
        parser = parser.add_flag("a", "")
                       .add_flag("d", "")
                       .add_opt("s", "")
                       .add_opt("f", "");
        parser.initialize(args.into_iter());
        assert!(parser.flagged(&'a'));
        assert!(!parser.flagged(&'d'));
        assert!(parser.get_opt(&'s') == Some(String::from("df")));
        assert!(parser.get_opt(&'f') == Some(String::from("foo")));
    }

    #[test]
    fn long_opts() {
        let args = vec![String::from("binname"), String::from("--foo=bar")];
        let mut parser = ArgParser::new(4);
        parser = parser.add_opt("", "foo");
        parser.initialize(args.into_iter());
        assert!(parser.get_opt("foo") == Some(String::from("bar")));
    }
}
<|endoftext|>


<commit_before>// Copyright 2015, Paul Osborne <osbpau@gmail.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/license/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option.  This file may not be copied, modified, or distributed
// except according to those terms.
//
// Portions of this implementation are based on work by Nat Pryce:
// https://github.com/npryce/rusty-pi/blob/master/src/pi/gpio.rs

#![crate_type = "lib"]
#![crate_name = "sysfs_gpio"]

///! GPIO access under Linux using the GPIO sysfs interface
///!
///! The methods exposed by this library are centered around
///! the `Pin` struct and map pretty directly the API exposed
///! by the kernel in syfs (https://www.kernel.org/doc/Documentation/gpio/sysfs.txt).
///!
///! # Examples
///!
///! Typical usage for systems where one wants to ensure that
///! the pins in use are unexported upon completion looks like
///! the follwoing:
///!
///! ```rust,ignore
///! extern crate sysfs_gpio;
///! use sysfs_gpio::Pin;
///!
///!
///! ```

extern crate nix;

use nix::sys::epoll::*;
use nix::unistd::close;

use std::io::prelude::*;
use std::os::unix::prelude::*;
use std::io;
use std::io::{Error, ErrorKind, SeekFrom};
use std::fs;
use std::fs::{File};

pub struct Pin {
    pin_num : u64,
}

#[derive(Clone,Debug)]
pub enum Direction {In, Out, High, Low}

#[derive(Clone,Debug)]
pub enum Edge {NoInterrupt, RisingEdge, FallingEdge, BothEdges}

#[macro_export]
macro_rules! try_unexport {
    ($gpio:ident, $e:expr) => (match $e {
        Ok(res) => res,
        Err(e) => { try!($gpio.unexport()); return Err(e) },
    });
}

fn from_nix_error(err: ::nix::Error) -> io::Error {
    io::Error::from_raw_os_error(err.errno() as i32)
}

/// Flush up to max bytes from the provided files input buffer
///
/// Typically, one would just use seek() for this sort of thing,
/// but for certain files (e.g. in sysfs), you need to actually
/// read it.
fn flush_input_from_file(dev_file: &mut File, max : usize) -> io::Result<usize> {
    let mut s = String::with_capacity(max);
    dev_file.read_to_string(&mut s)
}

/// Get the pin value from the provided file
fn get_value_from_file(dev_file: &mut File) -> io::Result<u8> {
    let mut s = String::with_capacity(10);
    try!(dev_file.seek(SeekFrom::Start(0)));
    try!(dev_file.read_to_string(&mut s));
    match s[..1].parse::<u8>() {
        Ok(n) => Ok(n),
        Err(_) => Err(Error::new(ErrorKind::Other, "Unexpected Error")),
    }
}

impl Pin {
    /// Write all of the provided contents to the specified devFile
    fn write_to_device_file(&self, dev_file_name: &str, value: &str) -> io::Result<()> {
        let gpio_path = format!("/sys/class/gpio/gpio{}/{}", self.pin_num, dev_file_name);
        let mut dev_file = try!(File::create(&gpio_path));
        try!(dev_file.write_all(value.as_bytes()));
        Ok(())
    }

    /// Create a new Pin with the provided `pin_num`
    ///
    /// This function does not export the provided pin_num.
    pub fn new(pin_num : u64) -> Pin {
        Pin {
            pin_num: pin_num,
        }
    }

    /// Run a closure with the GPIO exported
    ///
    /// Prior to the provided closure being executed, the Gpio
    /// will be eported.  After the closure execution is complete,
    /// the Gpio will be unexported.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// gpio = Pin::new(24);
    /// let res = gpio::with_exported(|| {
    ///     println!("At this point, the Pin is exported");
    ///     try!(gpio.set_direction(Direction::Low));
    ///     try!(gpio.set_value(1));
    ///     // ...
    /// };
    /// ```
    #[inline]
    pub fn with_exported<F: FnOnce() -> io::Result<()>>(&self, closure : F) -> io::Result<()> {
        try!(self.export());
        match closure() {
            Ok(()) => { try!(self.unexport()); Ok(()) },
            Err(err) => { try!(self.unexport()); Err(err) },
        }
    }

    /// Export the GPIO
    ///
    /// This is equivalent to `echo N > /sys/class/gpio/export` with
    /// the exception that the case where the GPIO is already exported
    /// is not an error.
    ///
    /// # Errors
    ///
    /// The main cases in which this function will fail and return an
    /// error are the following:
    /// 1. The system does not support the GPIO sysfs interface
    /// 2. The requested GPIO is out of range and cannot be exported
    /// 3. The requested GPIO is in use by the kernel and cannot
    ///    be exported by use in userspace
    ///
    /// # Example
    /// ```rust,ignore
    /// use sysfs_gpio::Pin;
    ///
    /// gpio = Pin::new(24);
    /// match gpio.export() {
    ///     Ok(()) => println!("Gpio {} exported!", gpio.pin),
    ///     Err(err) => println!("Gpio {} could not be exported: {}", gpio.pin, err),
    /// }
    /// ```
    pub fn export(&self) -> io::Result<()> {
        match fs::metadata(&format!("/sys/class/gpio/gpio{}", self.pin_num)) {
            Ok(_) => {},
            Err(_) => {
                let mut export_file = try!(File::create("/sys/class/gpio/export"));
                try!(export_file.write_all(format!("{}", self.pin_num).as_bytes()));
            }
        };
        Ok(())
    }

    /// Unexport the GPIO
    ///
    /// This function will unexport the provided by from syfs if
    /// it is currently exported.  If the pin is not currently
    /// exported, it will return without error.  That is, whenever
    /// this function returns Ok, the GPIO is not exported.
    pub fn unexport(&self) -> io::Result<()> {
        match fs::metadata(&format!("/sys/class/gpio/gpio{}", self.pin_num)) {
            Ok(_) => {
                let mut unexport_file = try!(File::create("/sys/class/gpio/unexport"));
                try!(unexport_file.write_all(format!("{}", self.pin_num).as_bytes()));
            },
            Err(_) => {} // not exported
        };
        Ok(())
    }

    /// Set this GPIO as either an input or an output
    ///
    /// The basic values allowed here are `Direction::In` and
    /// `Direction::Out` which set the Pin as either an input
    /// or output respectively.  In addition to those, two
    /// additional settings of `Direction::High` and
    /// `Direction::Low`.  These both set the Pin as an output
    /// but do so with an initial value of high or low respectively.
    /// This allows for glitch-free operation.
    ///
    /// Note that this entry may not exist if the kernel does
    /// not support changing the direction of a pin in userspace.  If
    /// this is the case, you will get an error.
    pub fn set_direction(&self, dir : Direction) -> io::Result<()> {
        self.write_to_device_file("direction", match dir {
            Direction::In => "in",
            Direction::Out => "out",
            Direction::High => "high",
            Direction::Low => "low",
        })
    }

    /// Get the value of the Pin (0 or 1)
    ///
    /// If successful, 1 will be returned if the pin is high
    /// and 0 will be returned if the pin is low (this may or may
    /// not match the signal level of the actual signal depending
    /// on the GPIO "active_low" entry).
    pub fn get_value(&self) -> io::Result<u8> {
        let mut dev_file = try!(File::open(&format!("/sys/class/gpio/gpio{}/value", self.pin_num)));
        get_value_from_file(&mut dev_file)
    }

    /// Set the value of the Pin
    ///
    /// This will set the value of the pin either high or low.
    /// A 0 value will set the pin low and any other value will
    /// set the pin high (1 is typical).
    pub fn set_value(&self, value : u8) -> io::Result<()> {
        let val = if value == 0 {
            "0"
        } else {
            "1"
        };
        self.write_to_device_file("value", val)
    }

    /// Set the edge on which this GPIO will trigger when polled
    ///
    /// The configured edge determines what changes to the Pin will
    /// result in `poll()` returning.  This call will return an Error
    /// if the pin does not allow interrupts.
    pub fn set_edge(&self, edge: Edge) -> io::Result<()> {
        self.write_to_device_file("edge", match edge {
            Edge::NoInterrupt => "none",
            Edge::RisingEdge => "rising",
            Edge::FallingEdge => "falling",
            Edge::BothEdges => "both",
        })
    }

    /// Get the currently configured edge for this pin
    ///
    /// This value will only be present if the Pin allows
    /// for interrupts.
    pub fn get_edge(&self) -> io::Result<Edge> {
        let mut dev_file = try!(File::open(&format!("/sys/class/gpio/gpio{}/edge", self.pin_num)));
        let mut s = String::with_capacity(10);
        try!(dev_file.read_to_string(&mut s));
        match s.trim() {
            "none" => Ok(Edge::NoInterrupt),
            "rising" => Ok(Edge::RisingEdge),
            "falling" => Ok(Edge::FallingEdge),
            "both" => Ok(Edge::BothEdges),
            other => Err(Error::new(ErrorKind::Other, format!("Unexpected edge file contents {}", other))),
        }
    }

    /// Get a PinPoller object for this pin
    ///
    /// This pin poller object will register an interrupt with the
    /// kernel and allow you to poll() on it and receive notifications
    /// that an interrupt has occured with minimal delay.
    pub fn get_poller(&self) -> io::Result<PinPoller> {
        PinPoller::new(self.pin_num)
    }
}

pub struct PinPoller {
    pin_num : u64,
    epoll_fd : RawFd,
    devfile : File,
}

impl PinPoller {

    /// Get the pin associated with this PinPoller
    ///
    /// Note that this will be a new Pin object with the
    /// proper pin number.
    pub fn get_pin(&self) -> Pin {
        return Pin::new(self.pin_num);
    }

    /// Create a new PinPoller for the provided pin number
    pub fn new(pin_num : u64) -> io::Result<PinPoller> {
        let devfile : File = try!(File::open(&format!("/sys/class/gpio/gpio{}/value", pin_num)));
        let devfile_fd = devfile.as_raw_fd();
        let epoll_fd = try!(epoll_create().map_err(from_nix_error));
        let events = EPOLLPRI | EPOLLET;
        let info = EpollEvent {
            events: events,
            data: 0u64,
        };

        match epoll_ctl(epoll_fd, EpollOp::EpollCtlAdd, devfile_fd, &info) {
            Ok(_) => {
                let pp = PinPoller {
                    pin_num: pin_num,
                    devfile: devfile,
                    epoll_fd: epoll_fd,
                };
                Ok(pp)
            },
            Err(err) => {
                let _ = close(epoll_fd); // cleanup
                Err(from_nix_error(err))
            }
        }
    }

    /// Block until an interrupt occurs
    ///
    /// This call will block until an interrupt occurs.  The types
    /// of interrupts which may result in this call returning
    /// may be configured by calling `set_edge()` prior to
    /// makeing this call.  This call makes use of epoll under the
    /// covers.  If it is desirable to poll on multiple GPIOs or
    /// other event source, you will need to implement that logic
    /// yourself.
    ///
    /// This function will return Some(value) of the pin if a change is
    /// detected or None if a timeout occurs.  Note that the value provided
    /// is the value of the pin as soon as we get to hanlding the interrupt
    /// in userspace.  Each time this fuction returns with a value, a change
    /// has occurred, but you could end up reading the same value multiple
    /// times as the value has changed back between when the interrupt
    /// ocurred and the current time.
    pub fn poll(&mut self, timeout_ms: isize) -> io::Result<Option<u8>> {
        try!(flush_input_from_file(&mut self.devfile, 255));
        let dummy_event = EpollEvent { events: EPOLLPRI | EPOLLET, data: 0u64};
        let mut events: [EpollEvent; 1] = [ dummy_event ];
        let cnt = try!(epoll_wait(self.epoll_fd, &mut events, timeout_ms).map_err(from_nix_error));
        Ok(match cnt {
            0 => None, // timeout
            _ => Some(try!(get_value_from_file(&mut self.devfile))),
        })
    }
}

impl Drop for PinPoller {
    fn drop(&mut self) {
        // we implement drop to close the underlying epoll fd as
        // it does not implement drop itself.  This is similar to
        // how mio works
        close(self.epoll_fd).unwrap();  // panic! if close files
    }
}
<commit_msg>Fixup top-of-module example<commit_after>// Copyright 2015, Paul Osborne <osbpau@gmail.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/license/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option.  This file may not be copied, modified, or distributed
// except according to those terms.
//
// Portions of this implementation are based on work by Nat Pryce:
// https://github.com/npryce/rusty-pi/blob/master/src/pi/gpio.rs

#![crate_type = "lib"]
#![crate_name = "sysfs_gpio"]

//! GPIO access under Linux using the GPIO sysfs interface
//!
//! The methods exposed by this library are centered around
//! the `Pin` struct and map pretty directly the API exposed
//! by the kernel in syfs (https://www.kernel.org/doc/Documentation/gpio/sysfs.txt).
//!
//! # Examples
//!
//! Typical usage for systems where one wants to ensure that
//! the pins in use are unexported upon completion looks like
//! the following:
//!
//! ```no_run
//! extern crate sysfs_gpio;
//!
//! use sysfs_gpio::{Direction, Pin};
//! use std::thread::sleep_ms;
//!
//! fn main() {
//!     let my_led = Pin::new(127); // number depends on chip, etc.
//!     my_led.with_exported(|| {
//!         loop {
//!             my_led.set_value(0).unwrap();
//!             sleep_ms(200);
//!             my_led.set_value(1).unwrap();
//!             sleep_ms(200);
//!         }
//!     }).unwrap();
//! }
//! ```

extern crate nix;

use nix::sys::epoll::*;
use nix::unistd::close;

use std::io::prelude::*;
use std::os::unix::prelude::*;
use std::io;
use std::io::{Error, ErrorKind, SeekFrom};
use std::fs;
use std::fs::{File};

pub struct Pin {
    pin_num : u64,
}

#[derive(Clone,Debug)]
pub enum Direction {In, Out, High, Low}

#[derive(Clone,Debug)]
pub enum Edge {NoInterrupt, RisingEdge, FallingEdge, BothEdges}

#[macro_export]
macro_rules! try_unexport {
    ($gpio:ident, $e:expr) => (match $e {
        Ok(res) => res,
        Err(e) => { try!($gpio.unexport()); return Err(e) },
    });
}

fn from_nix_error(err: ::nix::Error) -> io::Error {
    io::Error::from_raw_os_error(err.errno() as i32)
}

/// Flush up to max bytes from the provided files input buffer
///
/// Typically, one would just use seek() for this sort of thing,
/// but for certain files (e.g. in sysfs), you need to actually
/// read it.
fn flush_input_from_file(dev_file: &mut File, max : usize) -> io::Result<usize> {
    let mut s = String::with_capacity(max);
    dev_file.read_to_string(&mut s)
}

/// Get the pin value from the provided file
fn get_value_from_file(dev_file: &mut File) -> io::Result<u8> {
    let mut s = String::with_capacity(10);
    try!(dev_file.seek(SeekFrom::Start(0)));
    try!(dev_file.read_to_string(&mut s));
    match s[..1].parse::<u8>() {
        Ok(n) => Ok(n),
        Err(_) => Err(Error::new(ErrorKind::Other, "Unexpected Error")),
    }
}

impl Pin {
    /// Write all of the provided contents to the specified devFile
    fn write_to_device_file(&self, dev_file_name: &str, value: &str) -> io::Result<()> {
        let gpio_path = format!("/sys/class/gpio/gpio{}/{}", self.pin_num, dev_file_name);
        let mut dev_file = try!(File::create(&gpio_path));
        try!(dev_file.write_all(value.as_bytes()));
        Ok(())
    }

    /// Create a new Pin with the provided `pin_num`
    ///
    /// This function does not export the provided pin_num.
    pub fn new(pin_num : u64) -> Pin {
        Pin {
            pin_num: pin_num,
        }
    }

    /// Run a closure with the GPIO exported
    ///
    /// Prior to the provided closure being executed, the Gpio
    /// will be eported.  After the closure execution is complete,
    /// the Gpio will be unexported.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// gpio = Pin::new(24);
    /// let res = gpio::with_exported(|| {
    ///     println!("At this point, the Pin is exported");
    ///     try!(gpio.set_direction(Direction::Low));
    ///     try!(gpio.set_value(1));
    ///     // ...
    /// };
    /// ```
    #[inline]
    pub fn with_exported<F: FnOnce() -> io::Result<()>>(&self, closure : F) -> io::Result<()> {
        try!(self.export());
        match closure() {
            Ok(()) => { try!(self.unexport()); Ok(()) },
            Err(err) => { try!(self.unexport()); Err(err) },
        }
    }

    /// Export the GPIO
    ///
    /// This is equivalent to `echo N > /sys/class/gpio/export` with
    /// the exception that the case where the GPIO is already exported
    /// is not an error.
    ///
    /// # Errors
    ///
    /// The main cases in which this function will fail and return an
    /// error are the following:
    /// 1. The system does not support the GPIO sysfs interface
    /// 2. The requested GPIO is out of range and cannot be exported
    /// 3. The requested GPIO is in use by the kernel and cannot
    ///    be exported by use in userspace
    ///
    /// # Example
    /// ```rust,ignore
    /// use sysfs_gpio::Pin;
    ///
    /// gpio = Pin::new(24);
    /// match gpio.export() {
    ///     Ok(()) => println!("Gpio {} exported!", gpio.pin),
    ///     Err(err) => println!("Gpio {} could not be exported: {}", gpio.pin, err),
    /// }
    /// ```
    pub fn export(&self) -> io::Result<()> {
        match fs::metadata(&format!("/sys/class/gpio/gpio{}", self.pin_num)) {
            Ok(_) => {},
            Err(_) => {
                let mut export_file = try!(File::create("/sys/class/gpio/export"));
                try!(export_file.write_all(format!("{}", self.pin_num).as_bytes()));
            }
        };
        Ok(())
    }

    /// Unexport the GPIO
    ///
    /// This function will unexport the provided by from syfs if
    /// it is currently exported.  If the pin is not currently
    /// exported, it will return without error.  That is, whenever
    /// this function returns Ok, the GPIO is not exported.
    pub fn unexport(&self) -> io::Result<()> {
        match fs::metadata(&format!("/sys/class/gpio/gpio{}", self.pin_num)) {
            Ok(_) => {
                let mut unexport_file = try!(File::create("/sys/class/gpio/unexport"));
                try!(unexport_file.write_all(format!("{}", self.pin_num).as_bytes()));
            },
            Err(_) => {} // not exported
        };
        Ok(())
    }

    /// Set this GPIO as either an input or an output
    ///
    /// The basic values allowed here are `Direction::In` and
    /// `Direction::Out` which set the Pin as either an input
    /// or output respectively.  In addition to those, two
    /// additional settings of `Direction::High` and
    /// `Direction::Low`.  These both set the Pin as an output
    /// but do so with an initial value of high or low respectively.
    /// This allows for glitch-free operation.
    ///
    /// Note that this entry may not exist if the kernel does
    /// not support changing the direction of a pin in userspace.  If
    /// this is the case, you will get an error.
    pub fn set_direction(&self, dir : Direction) -> io::Result<()> {
        self.write_to_device_file("direction", match dir {
            Direction::In => "in",
            Direction::Out => "out",
            Direction::High => "high",
            Direction::Low => "low",
        })
    }

    /// Get the value of the Pin (0 or 1)
    ///
    /// If successful, 1 will be returned if the pin is high
    /// and 0 will be returned if the pin is low (this may or may
    /// not match the signal level of the actual signal depending
    /// on the GPIO "active_low" entry).
    pub fn get_value(&self) -> io::Result<u8> {
        let mut dev_file = try!(File::open(&format!("/sys/class/gpio/gpio{}/value", self.pin_num)));
        get_value_from_file(&mut dev_file)
    }

    /// Set the value of the Pin
    ///
    /// This will set the value of the pin either high or low.
    /// A 0 value will set the pin low and any other value will
    /// set the pin high (1 is typical).
    pub fn set_value(&self, value : u8) -> io::Result<()> {
        let val = if value == 0 {
            "0"
        } else {
            "1"
        };
        self.write_to_device_file("value", val)
    }

    /// Set the edge on which this GPIO will trigger when polled
    ///
    /// The configured edge determines what changes to the Pin will
    /// result in `poll()` returning.  This call will return an Error
    /// if the pin does not allow interrupts.
    pub fn set_edge(&self, edge: Edge) -> io::Result<()> {
        self.write_to_device_file("edge", match edge {
            Edge::NoInterrupt => "none",
            Edge::RisingEdge => "rising",
            Edge::FallingEdge => "falling",
            Edge::BothEdges => "both",
        })
    }

    /// Get the currently configured edge for this pin
    ///
    /// This value will only be present if the Pin allows
    /// for interrupts.
    pub fn get_edge(&self) -> io::Result<Edge> {
        let mut dev_file = try!(File::open(&format!("/sys/class/gpio/gpio{}/edge", self.pin_num)));
        let mut s = String::with_capacity(10);
        try!(dev_file.read_to_string(&mut s));
        match s.trim() {
            "none" => Ok(Edge::NoInterrupt),
            "rising" => Ok(Edge::RisingEdge),
            "falling" => Ok(Edge::FallingEdge),
            "both" => Ok(Edge::BothEdges),
            other => Err(Error::new(ErrorKind::Other, format!("Unexpected edge file contents {}", other))),
        }
    }

    /// Get a PinPoller object for this pin
    ///
    /// This pin poller object will register an interrupt with the
    /// kernel and allow you to poll() on it and receive notifications
    /// that an interrupt has occured with minimal delay.
    pub fn get_poller(&self) -> io::Result<PinPoller> {
        PinPoller::new(self.pin_num)
    }
}

pub struct PinPoller {
    pin_num : u64,
    epoll_fd : RawFd,
    devfile : File,
}

impl PinPoller {

    /// Get the pin associated with this PinPoller
    ///
    /// Note that this will be a new Pin object with the
    /// proper pin number.
    pub fn get_pin(&self) -> Pin {
        return Pin::new(self.pin_num);
    }

    /// Create a new PinPoller for the provided pin number
    pub fn new(pin_num : u64) -> io::Result<PinPoller> {
        let devfile : File = try!(File::open(&format!("/sys/class/gpio/gpio{}/value", pin_num)));
        let devfile_fd = devfile.as_raw_fd();
        let epoll_fd = try!(epoll_create().map_err(from_nix_error));
        let events = EPOLLPRI | EPOLLET;
        let info = EpollEvent {
            events: events,
            data: 0u64,
        };

        match epoll_ctl(epoll_fd, EpollOp::EpollCtlAdd, devfile_fd, &info) {
            Ok(_) => {
                let pp = PinPoller {
                    pin_num: pin_num,
                    devfile: devfile,
                    epoll_fd: epoll_fd,
                };
                Ok(pp)
            },
            Err(err) => {
                let _ = close(epoll_fd); // cleanup
                Err(from_nix_error(err))
            }
        }
    }

    /// Block until an interrupt occurs
    ///
    /// This call will block until an interrupt occurs.  The types
    /// of interrupts which may result in this call returning
    /// may be configured by calling `set_edge()` prior to
    /// makeing this call.  This call makes use of epoll under the
    /// covers.  If it is desirable to poll on multiple GPIOs or
    /// other event source, you will need to implement that logic
    /// yourself.
    ///
    /// This function will return Some(value) of the pin if a change is
    /// detected or None if a timeout occurs.  Note that the value provided
    /// is the value of the pin as soon as we get to hanlding the interrupt
    /// in userspace.  Each time this fuction returns with a value, a change
    /// has occurred, but you could end up reading the same value multiple
    /// times as the value has changed back between when the interrupt
    /// ocurred and the current time.
    pub fn poll(&mut self, timeout_ms: isize) -> io::Result<Option<u8>> {
        try!(flush_input_from_file(&mut self.devfile, 255));
        let dummy_event = EpollEvent { events: EPOLLPRI | EPOLLET, data: 0u64};
        let mut events: [EpollEvent; 1] = [ dummy_event ];
        let cnt = try!(epoll_wait(self.epoll_fd, &mut events, timeout_ms).map_err(from_nix_error));
        Ok(match cnt {
            0 => None, // timeout
            _ => Some(try!(get_value_from_file(&mut self.devfile))),
        })
    }
}

impl Drop for PinPoller {
    fn drop(&mut self) {
        // we implement drop to close the underlying epoll fd as
        // it does not implement drop itself.  This is similar to
        // how mio works
        close(self.epoll_fd).unwrap();  // panic! if close files
    }
}
<|endoftext|>


<commit_before>//! Use this library to open a path or URL using the program configured on the system.
//!
//! # Usage
//!
//! Open the given URL in the default web browser.
//!
//! ```no_run
//! open::that("http://rust-lang.org").unwrap();
//! ```
//!
//! Alternatively, specify the program to be used to open the path or URL.
//!
//! ```no_run
//! open::with("http://rust-lang.org", "firefox").unwrap();
//! ```
//!
//! # Notes
//!
//! As an operating system program is used, the open operation can fail.
//! Therefore, you are advised to at least check the result and behave
//! accordingly, e.g. by letting the user know that the open operation failed.
//!
//! ```no_run
//! let path = "http://rust-lang.org";
//!
//! match open::that(path) {
//!     Ok(()) => println!("Opened '{}' successfully.", path),
//!     Err(err) => eprintln!("An error occurred when opening '{}': {}", path, err),
//! }
//! ```

#[cfg(target_os = "windows")]
use windows as os;

#[cfg(target_os = "macos")]
use macos as os;

#[cfg(target_os = "ios")]
use ios as os;

#[cfg(any(
    target_os = "linux",
    target_os = "android",
    target_os = "freebsd",
    target_os = "dragonfly",
    target_os = "netbsd",
    target_os = "openbsd",
    target_os = "solaris"
))]
use unix as os;

#[cfg(not(any(
    target_os = "linux",
    target_os = "android",
    target_os = "freebsd",
    target_os = "dragonfly",
    target_os = "netbsd",
    target_os = "openbsd",
    target_os = "solaris",
    target_os = "ios",
    target_os = "macos",
    target_os = "windows",
)))]
compile_error!("open is not supported on this platform");

use std::{
    ffi::OsStr,
    io,
    process::{Command, Output, Stdio},
    thread,
};

type Result = io::Result<()>;

/// Open path with the default application.
///
/// # Examples
///
/// ```no_run
/// let path = "http://rust-lang.org";
///
/// match open::that(path) {
///     Ok(()) => println!("Opened '{}' successfully.", path),
///     Err(err) => panic!("An error occurred when opening '{}': {}", path, err),
/// }
/// ```
///
/// # Errors
///
/// A [`std::io::Error`] is returned on failure. Because different operating systems
/// handle errors differently it is recommend to not match on a certain error.
pub fn that<T: AsRef<OsStr> + Sized>(path: T) -> Result {
    os::that(path)
}

/// Open path with the given application.
///
/// # Examples
///
/// ```no_run
/// let path = "http://rust-lang.org";
/// let app = "firefox";
///
/// match open::with(path, app) {
///     Ok(()) => println!("Opened '{}' successfully.", path),
///     Err(err) => panic!("An error occurred when opening '{}': {}", path, err),
/// }
/// ```
///
/// # Errors
///
/// A [`std::io::Error`] is returned on failure. Because different operating systems
/// handle errors differently it is recommend to not match on a certain error.
pub fn with<T: AsRef<OsStr> + Sized>(path: T, app: impl Into<String>) -> Result {
    os::with(path, app)
}

/// Open path with the default application in a new thread.
///
/// See documentation of [`that`] for more details.
pub fn that_in_background<T: AsRef<OsStr> + Sized>(path: T) -> thread::JoinHandle<Result> {
    let path = path.as_ref().to_os_string();
    thread::spawn(|| that(path))
}

/// Open path with the given application in a new thread.
///
/// See documentation of [`with`] for more details.
pub fn with_in_background<T: AsRef<OsStr> + Sized>(
    path: T,
    app: impl Into<String>,
) -> thread::JoinHandle<Result> {
    let path = path.as_ref().to_os_string();
    let app = app.into();
    thread::spawn(|| with(path, app))
}

trait IntoResult<T> {
    fn into_result(self) -> T;
}

impl IntoResult<Result> for io::Result<Output> {
    fn into_result(self) -> Result {
        match self {
            Ok(o) if o.status.success() => Ok(()),
            Ok(o) => Err(from_output(o)),
            Err(err) => Err(err),
        }
    }
}

#[cfg(windows)]
impl IntoResult<Result> for winapi::ctypes::c_int {
    fn into_result(self) -> Result {
        match self {
            i if i > 32 => Ok(()),
            _ => Err(io::Error::last_os_error()),
        }
    }
}

fn from_output(output: Output) -> io::Error {
    let error_msg = match output.stderr.is_empty() {
        true => output.status.to_string(),
        false => format!(
            "{} ({})",
            String::from_utf8_lossy(&output.stderr).trim(),
            output.status
        ),
    };

    io::Error::new(io::ErrorKind::Other, error_msg)
}

trait CommandExt {
    fn output_stderr(&mut self) -> io::Result<Output>;
}

impl CommandExt for Command {
    fn output_stderr(&mut self) -> io::Result<Output> {
        let mut process = self
            .stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::piped())
            .spawn()?;

        let status = process.wait()?;

        // Read up to 256 bytes from stderr.
        use std::io::Read;
        let mut stderr = vec![0; 256];
        let len = process
            .stderr
            .take()
            .and_then(|mut err| err.read(&mut stderr).ok())
            .unwrap_or(0);
        stderr.truncate(len);

        Ok(Output {
            status,
            stderr,
            stdout: vec![],
        })
    }
}

#[cfg(windows)]
mod windows {
    use std::{ffi::OsStr, io, os::windows::ffi::OsStrExt, ptr};

    use winapi::ctypes::c_int;
    use winapi::um::shellapi::ShellExecuteW;

    use crate::{IntoResult, Result};

    fn convert_path(path: &OsStr) -> io::Result<Vec<u16>> {
        let mut maybe_result: Vec<_> = path.encode_wide().collect();
        if maybe_result.iter().any(|&u| u == 0) {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "path contains NUL byte(s)",
            ));
        }
        maybe_result.push(0);
        Ok(maybe_result)
    }

    pub fn that<T: AsRef<OsStr> + Sized>(path: T) -> Result {
        const SW_SHOW: c_int = 5;

        let path = convert_path(path.as_ref())?;
        let operation: Vec<u16> = OsStr::new("open\0").encode_wide().collect();
        let result = unsafe {
            ShellExecuteW(
                ptr::null_mut(),
                operation.as_ptr(),
                path.as_ptr(),
                ptr::null(),
                ptr::null(),
                SW_SHOW,
            )
        };
        (result as c_int).into_result()
    }

    pub fn with<T: AsRef<OsStr> + Sized>(path: T, app: impl Into<String>) -> Result {
        const SW_SHOW: c_int = 5;

        let path = convert_path(path.as_ref())?;
        let operation: Vec<u16> = OsStr::new("open\0").encode_wide().collect();
        let app_name: Vec<u16> = OsStr::new(&format!("{}\0", app.into()))
            .encode_wide()
            .collect();
        let result = unsafe {
            ShellExecuteW(
                ptr::null_mut(),
                operation.as_ptr(),
                app_name.as_ptr(),
                path.as_ptr(),
                ptr::null(),
                SW_SHOW,
            )
        };
        (result as c_int).into_result()
    }
}

#[cfg(target_os = "macos")]
mod macos {
    use std::{ffi::OsStr, process::Command};

    use crate::{CommandExt, IntoResult, Result};

    pub fn that<T: AsRef<OsStr> + Sized>(path: T) -> Result {
        Command::new("open")
            .arg(path.as_ref())
            .output_stderr()
            .into_result()
    }

    pub fn with<T: AsRef<OsStr> + Sized>(path: T, app: impl Into<String>) -> Result {
        Command::new("open")
            .arg(path.as_ref())
            .arg("-a")
            .arg(app.into())
            .output_stderr()
            .into_result()
    }
}

#[cfg(target_os = "ios")]
mod ios {
    use std::{ffi::OsStr, process::Command};

    use crate::{CommandExt, IntoResult, Result};

    pub fn that<T: AsRef<OsStr> + Sized>(path: T) -> Result {
        Command::new("uiopen")
            .arg("--url")
            .arg(path.as_ref())
            .output_stderr()
            .into_result()
    }

    pub fn with<T: AsRef<OsStr> + Sized>(path: T, app: impl Into<String>) -> Result {
        Command::new("uiopen")
            .arg("--url")
            .arg(path.as_ref())
            .arg("--bundleid")
            .arg(app.into())
            .output_stderr()
            .into_result()
    }
}

#[cfg(any(
    target_os = "linux",
    target_os = "android",
    target_os = "freebsd",
    target_os = "dragonfly",
    target_os = "netbsd",
    target_os = "openbsd",
    target_os = "solaris"
))]
mod unix {
    use std::{
        env,
        ffi::{OsStr, OsString},
        path::{Path, PathBuf},
        process::Command,
    };

    use crate::{CommandExt, IntoResult, Result};

    pub fn that<T: AsRef<OsStr> + Sized>(path: T) -> Result {
        let path = path.as_ref();
        let open_handlers = [
            ("xdg-open", &[path] as &[_]),
            ("gio", &[OsStr::new("open"), path]),
            ("gnome-open", &[path]),
            ("kde-open", &[path]),
            ("wslview", &[&wsl_path(path)]),
        ];

        let mut unsuccessful = None;
        let mut io_error = None;

        for (command, args) in &open_handlers {
            let result = Command::new(command).args(*args).output_stderr();

            match result {
                Ok(o) if o.status.success() => return Ok(()),
                Ok(o) => unsuccessful = unsuccessful.or_else(|| Some(crate::from_output(o))),
                Err(err) => io_error = io_error.or(Some(err)),
            }
        }

        Err(unsuccessful
            .or(io_error)
            .expect("successful cases don't get here"))
    }

    pub fn with<T: AsRef<OsStr> + Sized>(path: T, app: impl Into<String>) -> Result {
        Command::new(app.into())
            .arg(path.as_ref())
            .output_stderr()
            .into_result()
    }

    // Polyfill to workaround absolute path bug in wslu(wslview). In versions before
    // v3.1.1, wslview is unable to find absolute paths. `wsl_path` converts an
    // absolute path into a relative path starting from the current directory. If
    // the path is already a relative path or the conversion fails the original path
    // is returned.
    fn wsl_path<T: AsRef<OsStr>>(path: T) -> OsString {
        fn path_relative_to_current_dir<T: AsRef<OsStr>>(path: T) -> Option<PathBuf> {
            let path = Path::new(&path);

            if path.is_relative() {
                return None;
            }

            let base = env::current_dir().ok()?;
            pathdiff::diff_paths(path, base)
        }

        match path_relative_to_current_dir(&path) {
            None => OsString::from(&path),
            Some(relative) => OsString::from(relative),
        }
    }
}
<commit_msg>fix: Prevent deadlocks due to filled pipe on stderr (#85)<commit_after>//! Use this library to open a path or URL using the program configured on the system.
//!
//! # Usage
//!
//! Open the given URL in the default web browser.
//!
//! ```no_run
//! open::that("http://rust-lang.org").unwrap();
//! ```
//!
//! Alternatively, specify the program to be used to open the path or URL.
//!
//! ```no_run
//! open::with("http://rust-lang.org", "firefox").unwrap();
//! ```
//!
//! # Notes
//!
//! As an operating system program is used, the open operation can fail.
//! Therefore, you are advised to at least check the result and behave
//! accordingly, e.g. by letting the user know that the open operation failed.
//!
//! ```no_run
//! let path = "http://rust-lang.org";
//!
//! match open::that(path) {
//!     Ok(()) => println!("Opened '{}' successfully.", path),
//!     Err(err) => eprintln!("An error occurred when opening '{}': {}", path, err),
//! }
//! ```

#[cfg(target_os = "windows")]
use windows as os;

#[cfg(target_os = "macos")]
use macos as os;

#[cfg(target_os = "ios")]
use ios as os;

#[cfg(any(
    target_os = "linux",
    target_os = "android",
    target_os = "freebsd",
    target_os = "dragonfly",
    target_os = "netbsd",
    target_os = "openbsd",
    target_os = "solaris"
))]
use unix as os;

#[cfg(not(any(
    target_os = "linux",
    target_os = "android",
    target_os = "freebsd",
    target_os = "dragonfly",
    target_os = "netbsd",
    target_os = "openbsd",
    target_os = "solaris",
    target_os = "ios",
    target_os = "macos",
    target_os = "windows",
)))]
compile_error!("open is not supported on this platform");

use std::{
    ffi::OsStr,
    io,
    process::{Command, Output, Stdio},
    thread,
};

type Result = io::Result<()>;

/// Open path with the default application.
///
/// # Examples
///
/// ```no_run
/// let path = "http://rust-lang.org";
///
/// match open::that(path) {
///     Ok(()) => println!("Opened '{}' successfully.", path),
///     Err(err) => panic!("An error occurred when opening '{}': {}", path, err),
/// }
/// ```
///
/// # Errors
///
/// A [`std::io::Error`] is returned on failure. Because different operating systems
/// handle errors differently it is recommend to not match on a certain error.
pub fn that<T: AsRef<OsStr> + Sized>(path: T) -> Result {
    os::that(path)
}

/// Open path with the given application.
///
/// # Examples
///
/// ```no_run
/// let path = "http://rust-lang.org";
/// let app = "firefox";
///
/// match open::with(path, app) {
///     Ok(()) => println!("Opened '{}' successfully.", path),
///     Err(err) => panic!("An error occurred when opening '{}': {}", path, err),
/// }
/// ```
///
/// # Errors
///
/// A [`std::io::Error`] is returned on failure. Because different operating systems
/// handle errors differently it is recommend to not match on a certain error.
pub fn with<T: AsRef<OsStr> + Sized>(path: T, app: impl Into<String>) -> Result {
    os::with(path, app)
}

/// Open path with the default application in a new thread.
///
/// See documentation of [`that`] for more details.
pub fn that_in_background<T: AsRef<OsStr> + Sized>(path: T) -> thread::JoinHandle<Result> {
    let path = path.as_ref().to_os_string();
    thread::spawn(|| that(path))
}

/// Open path with the given application in a new thread.
///
/// See documentation of [`with`] for more details.
pub fn with_in_background<T: AsRef<OsStr> + Sized>(
    path: T,
    app: impl Into<String>,
) -> thread::JoinHandle<Result> {
    let path = path.as_ref().to_os_string();
    let app = app.into();
    thread::spawn(|| with(path, app))
}

trait IntoResult<T> {
    fn into_result(self) -> T;
}

impl IntoResult<Result> for io::Result<Output> {
    fn into_result(self) -> Result {
        match self {
            Ok(o) if o.status.success() => Ok(()),
            Ok(o) => Err(from_output(o)),
            Err(err) => Err(err),
        }
    }
}

#[cfg(windows)]
impl IntoResult<Result> for winapi::ctypes::c_int {
    fn into_result(self) -> Result {
        match self {
            i if i > 32 => Ok(()),
            _ => Err(io::Error::last_os_error()),
        }
    }
}

fn from_output(output: Output) -> io::Error {
    let error_msg = match output.stderr.is_empty() {
        true => output.status.to_string(),
        false => format!(
            "{} ({})",
            String::from_utf8_lossy(&output.stderr).trim(),
            output.status
        ),
    };

    io::Error::new(io::ErrorKind::Other, error_msg)
}

trait CommandExt {
    fn output_stderr(&mut self) -> io::Result<Output>;
}

impl CommandExt for Command {
    fn output_stderr(&mut self) -> io::Result<Output> {
        let mut process = self
            .stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::piped())
            .spawn()?;

        // Consume all stderr - it's open just for a few programs which can't handle it being closed.
        use std::io::Read;
        let mut stderr = vec![0; 256];
        let mut stderr_src = process.stderr.take().expect("piped stderr");

        let len = stderr_src.read(&mut stderr).unwrap_or(0);
        stderr.truncate(len);

        // consume the rest to avoid blocking
        std::io::copy(&mut stderr_src, &mut std::io::sink()).ok();

        let status = process.wait()?;
        Ok(Output {
            status,
            stderr,
            stdout: vec![],
        })
    }
}

#[cfg(windows)]
mod windows {
    use std::{ffi::OsStr, io, os::windows::ffi::OsStrExt, ptr};

    use winapi::ctypes::c_int;
    use winapi::um::shellapi::ShellExecuteW;

    use crate::{IntoResult, Result};

    fn convert_path(path: &OsStr) -> io::Result<Vec<u16>> {
        let mut maybe_result: Vec<_> = path.encode_wide().collect();
        if maybe_result.iter().any(|&u| u == 0) {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "path contains NUL byte(s)",
            ));
        }
        maybe_result.push(0);
        Ok(maybe_result)
    }

    pub fn that<T: AsRef<OsStr> + Sized>(path: T) -> Result {
        const SW_SHOW: c_int = 5;

        let path = convert_path(path.as_ref())?;
        let operation: Vec<u16> = OsStr::new("open\0").encode_wide().collect();
        let result = unsafe {
            ShellExecuteW(
                ptr::null_mut(),
                operation.as_ptr(),
                path.as_ptr(),
                ptr::null(),
                ptr::null(),
                SW_SHOW,
            )
        };
        (result as c_int).into_result()
    }

    pub fn with<T: AsRef<OsStr> + Sized>(path: T, app: impl Into<String>) -> Result {
        const SW_SHOW: c_int = 5;

        let path = convert_path(path.as_ref())?;
        let operation: Vec<u16> = OsStr::new("open\0").encode_wide().collect();
        let app_name: Vec<u16> = OsStr::new(&format!("{}\0", app.into()))
            .encode_wide()
            .collect();
        let result = unsafe {
            ShellExecuteW(
                ptr::null_mut(),
                operation.as_ptr(),
                app_name.as_ptr(),
                path.as_ptr(),
                ptr::null(),
                SW_SHOW,
            )
        };
        (result as c_int).into_result()
    }
}

#[cfg(target_os = "macos")]
mod macos {
    use std::{ffi::OsStr, process::Command};

    use crate::{CommandExt, IntoResult, Result};

    pub fn that<T: AsRef<OsStr> + Sized>(path: T) -> Result {
        Command::new("open")
            .arg(path.as_ref())
            .output_stderr()
            .into_result()
    }

    pub fn with<T: AsRef<OsStr> + Sized>(path: T, app: impl Into<String>) -> Result {
        Command::new("open")
            .arg(path.as_ref())
            .arg("-a")
            .arg(app.into())
            .output_stderr()
            .into_result()
    }
}

#[cfg(target_os = "ios")]
mod ios {
    use std::{ffi::OsStr, process::Command};

    use crate::{CommandExt, IntoResult, Result};

    pub fn that<T: AsRef<OsStr> + Sized>(path: T) -> Result {
        Command::new("uiopen")
            .arg("--url")
            .arg(path.as_ref())
            .output_stderr()
            .into_result()
    }

    pub fn with<T: AsRef<OsStr> + Sized>(path: T, app: impl Into<String>) -> Result {
        Command::new("uiopen")
            .arg("--url")
            .arg(path.as_ref())
            .arg("--bundleid")
            .arg(app.into())
            .output_stderr()
            .into_result()
    }
}

#[cfg(any(
    target_os = "linux",
    target_os = "android",
    target_os = "freebsd",
    target_os = "dragonfly",
    target_os = "netbsd",
    target_os = "openbsd",
    target_os = "solaris"
))]
mod unix {
    use std::{
        env,
        ffi::{OsStr, OsString},
        path::{Path, PathBuf},
        process::Command,
    };

    use crate::{CommandExt, IntoResult, Result};

    pub fn that<T: AsRef<OsStr> + Sized>(path: T) -> Result {
        let path = path.as_ref();
        let open_handlers = [
            ("xdg-open", &[path] as &[_]),
            ("gio", &[OsStr::new("open"), path]),
            ("gnome-open", &[path]),
            ("kde-open", &[path]),
            ("wslview", &[&wsl_path(path)]),
        ];

        let mut unsuccessful = None;
        let mut io_error = None;

        for (command, args) in &open_handlers {
            let result = Command::new(command).args(*args).output_stderr();

            match result {
                Ok(o) if o.status.success() => return Ok(()),
                Ok(o) => unsuccessful = unsuccessful.or_else(|| Some(crate::from_output(o))),
                Err(err) => io_error = io_error.or(Some(err)),
            }
        }

        Err(unsuccessful
            .or(io_error)
            .expect("successful cases don't get here"))
    }

    pub fn with<T: AsRef<OsStr> + Sized>(path: T, app: impl Into<String>) -> Result {
        Command::new(app.into())
            .arg(path.as_ref())
            .output_stderr()
            .into_result()
    }

    // Polyfill to workaround absolute path bug in wslu(wslview). In versions before
    // v3.1.1, wslview is unable to find absolute paths. `wsl_path` converts an
    // absolute path into a relative path starting from the current directory. If
    // the path is already a relative path or the conversion fails the original path
    // is returned.
    fn wsl_path<T: AsRef<OsStr>>(path: T) -> OsString {
        fn path_relative_to_current_dir<T: AsRef<OsStr>>(path: T) -> Option<PathBuf> {
            let path = Path::new(&path);

            if path.is_relative() {
                return None;
            }

            let base = env::current_dir().ok()?;
            pathdiff::diff_paths(path, base)
        }

        match path_relative_to_current_dir(&path) {
            None => OsString::from(&path),
            Some(relative) => OsString::from(relative),
        }
    }
}
<|endoftext|>