start sketching out some things

5 files changed, 673 insertions, 0 deletions

diff --git a/src/colors.rs b/src/colors.rs
new file mode 100644
index 0000000..09492e3
--- /dev/null
+++ b/src/colors.rs
@@ -0,0 +1,126 @@
+use std;
+use term;
+
+#[derive(Debug,Clone)]
+pub enum ShellType {
+    Unknown,
+    Bash,
+    Zsh,
+}
+
+pub struct Colors {
+    color_map: std::collections::HashMap<String, term::color::Color>,
+    unknown_color: term::color::Color,
+    shell_type: ShellType,
+}
+
+impl ShellType {
+    pub fn from_str(shell: &str) -> Self {
+        match shell {
+            "bash" => ShellType::Bash,
+            "zsh" => ShellType::Zsh,
+            _ => panic!("unknown shell {}", shell)
+        }
+    }
+}
+
+impl Colors {
+    pub fn new(shell_type: ShellType) -> Colors {
+        let mut color_map = std::collections::HashMap::new();
+
+        color_map.insert("battery_warn".to_string(), term::color::YELLOW);
+        color_map.insert("battery_crit".to_string(), term::color::RED);
+        color_map.insert("battery_emerg".to_string(), term::color::BRIGHT_RED);
+        color_map.insert("battery_full".to_string(), term::color::GREEN);
+        color_map.insert("battery_charging".to_string(), term::color::GREEN);
+
+        color_map.insert("default".to_string(), term::color::BRIGHT_BLACK);
+        color_map.insert("error".to_string(), term::color::RED);
+
+        let unknown_color = term::color::YELLOW;
+
+        Colors {
+            color_map: color_map,
+            unknown_color: unknown_color,
+            shell_type: shell_type,
+        }
+    }
+
+    pub fn print(&self, color: &str, text: &str) {
+        let color = self.color_map.get(color);
+        self.print_with_color(color, text);
+    }
+
+    pub fn pad(&self, len: usize) {
+        print!("{}", " ".repeat(len));
+    }
+
+    pub fn newline(&self) {
+        self.print_wrapped(|| {
+            print!("{}", "\n");
+        });
+    }
+
+    pub fn print_host(&self, host: &Option<String>, text: &str) {
+        let color = host
+            .clone()
+            .and_then(|hostname| {
+                self.color_map.get(&format!("host_{}", hostname))
+            });
+        self.print_with_color(color, text);
+    }
+
+    pub fn print_user(&self, user: &Option<String>, text: &str) {
+        let color = user
+            .clone()
+            .and_then(|username| {
+                self.color_map.get(&format!("user_{}", username))
+            });
+        self.print_with_color(color, text);
+    }
+
+    fn print_with_color(&self, color: Option<&term::color::Color>, text: &str) {
+        let mut t = term::stdout().unwrap();
+        self.print_color(&mut *t, color);
+        write!(t, "{}", text).unwrap();
+        t.reset().unwrap();
+    }
+
+    fn print_color(&self, t: &mut term::StdoutTerminal, color: Option<&term::color::Color>) {
+        self.print_wrapped(|| {
+            let real_color = *color.unwrap_or(&self.unknown_color);
+            t.fg(real_color).unwrap();
+            match real_color {
+                term::color::BRIGHT_BLACK
+            | term::color::BRIGHT_BLUE
+            | term::color::BRIGHT_CYAN
+            | term::color::BRIGHT_GREEN
+            | term::color::BRIGHT_MAGENTA
+            | term::color::BRIGHT_RED
+            | term::color::BRIGHT_WHITE
+            | term::color::BRIGHT_YELLOW => {
+                    t.attr(term::Attr::Bold).unwrap()
+                },
+                _ => {},
+            }
+        })
+    }
+
+    fn print_wrapped<T>(&self, printer: T)
+        where T: FnOnce()
+    {
+        match self.shell_type {
+            ShellType::Bash => { print!("{}", "\\["); },
+            ShellType::Zsh => { print!("{}", "%{"); },
+            _ => {},
+        }
+
+        printer();
+
+        match self.shell_type {
+            ShellType::Bash => { print!("{}", "\\]"); },
+            ShellType::Zsh => { print!("{}", "%}"); },
+            _ => {},
+        }
+    }
+}
diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..471a3a2
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,54 @@
+extern crate chrono;
+extern crate clap;
+extern crate hostname;
+extern crate regex;
+extern crate term;
+extern crate term_size;
+extern crate users;
+extern crate walkdir;
+
+mod colors;
+mod power;
+mod prompt;
+mod system_info;
+
+fn collect_data() -> prompt::PromptData {
+    let matches = clap::App::new("fancy-prompt")
+        .about("Prints a fancy prompt")
+        // XXX author, version (extract from cargo)
+        .arg(clap::Arg::with_name("prompt-escape")
+             .long("prompt-escape")
+             .value_name("SHELL")
+             .help("Produces escape sequence wrappers for the given shell")
+             .takes_value(true))
+        .arg(clap::Arg::with_name("error-code")
+             .value_name("ERROR_CODE")
+             .help("The error code of the previously run command")
+             )
+        .get_matches();
+
+    prompt::PromptData {
+        shell: matches
+            .value_of("prompt-escape")
+            .map(|shell| colors::ShellType::from_str(shell))
+            .unwrap_or(colors::ShellType::Unknown),
+        error_code: matches
+            .value_of("error-code")
+            .map(|code| code.parse().expect("error code must be a u8"))
+            .unwrap_or(0),
+
+        hostname: system_info::hostname(),
+        terminal_cols: system_info::terminal_cols(),
+        pwd: system_info::pwd(),
+        home: system_info::home(),
+        user: system_info::user(),
+        is_root: system_info::is_root(),
+        time: system_info::time(),
+        power_info: system_info::power_info(),
+    }
+}
+
+fn main() {
+    let data = collect_data();
+    prompt::Prompt::new(data).display();
+}
diff --git a/src/power.rs b/src/power.rs
new file mode 100644
index 0000000..2a9dd31
--- /dev/null
+++ b/src/power.rs
@@ -0,0 +1,137 @@
+use std;
+use walkdir;
+
+use std::io::Read;
+
+// XXX maybe extract this out into a separate crate?
+
+#[derive(PartialEq,Eq,Debug,Clone)]
+enum PowerSupplyType {
+    AC,
+    Battery,
+}
+
+#[derive(Debug)]
+pub struct PowerInfo {
+    power_supplies: Vec<PowerSupplyInfo>,
+}
+
+#[derive(Debug,Clone)]
+struct PowerSupplyInfo {
+    name: String,
+    ty: PowerSupplyType,
+    energy_now: Option<u64>,
+    energy_full: Option<u64>,
+    online: Option<bool>,
+}
+
+impl PowerInfo {
+    pub fn new() -> PowerInfo {
+        let mut power_supplies = vec![];
+        for entry in walkdir::WalkDir::new("/sys/class/power_supply/").min_depth(1).max_depth(1).follow_links(true) {
+            let entry = entry.unwrap();
+
+            let name = entry.path()
+                .file_name()
+                .unwrap()
+                .to_string_lossy()
+                .into_owned();
+            let ty = slurp(entry.path().join("type"))
+                .map(|t: String| PowerSupplyType::from_str(&t))
+                .expect("couldn't find power supply type");
+            let full = slurp(entry.path().join("energy_full"));
+            let now = slurp(entry.path().join("energy_now"));
+            let online = slurp(entry.path().join("online"))
+                .map(|n: u8| n != 0);
+
+            power_supplies.push(
+                PowerSupplyInfo {
+                    name: name,
+                    ty: ty,
+                    energy_now: now,
+                    energy_full: full,
+                    online: online,
+                }
+            )
+        }
+
+        PowerInfo {
+            power_supplies: power_supplies,
+        }
+    }
+
+    pub fn battery_usage(&self) -> Option<f64> {
+        let mut total_now = 0;
+        let mut total_full = 0;
+        for battery in self.batteries() {
+            if let Some(now) = battery.energy_now {
+                total_now += now;
+            }
+            else {
+                return None;
+            }
+            if let Some(full) = battery.energy_full {
+                total_full += full;
+            }
+            else {
+                return None;
+            }
+        }
+
+        if total_full > 0 {
+            Some((total_now as f64) / (total_full as f64))
+        }
+        else {
+            None
+        }
+    }
+
+    pub fn charging(&self) -> bool {
+        for mains in self.mains() {
+            if mains.online == Some(true) {
+                return true;
+            }
+        }
+        false
+    }
+
+    fn batteries(&self) -> Vec<PowerSupplyInfo> {
+        self.power_supplies
+            .iter()
+            .cloned()
+            .filter(|p| p.ty == PowerSupplyType::Battery)
+            .collect()
+    }
+
+    fn mains(&self) -> Vec<PowerSupplyInfo> {
+        self.power_supplies
+            .iter()
+            .filter(|p| p.ty == PowerSupplyType::AC)
+            .cloned()
+            .collect()
+    }
+}
+
+impl PowerSupplyType {
+    fn from_str(ty: &str) -> Self {
+        match ty {
+            "Mains" => PowerSupplyType::AC,
+            "Battery" => PowerSupplyType::Battery,
+            _ => panic!("unknown power supply type {}", ty)
+        }
+    }
+}
+
+fn slurp<T, U>(path: U) -> Option<T>
+    where T: std::str::FromStr,
+          U: AsRef<std::path::Path>
+{
+    let mut contents = String::new();
+    std::fs::File::open(path)
+        .ok()
+        .and_then(|mut fh| {
+            fh.read_to_string(&mut contents)
+                .ok()
+                .and_then(|_| contents.trim().parse().ok())
+        })
+}
diff --git a/src/prompt.rs b/src/prompt.rs
new file mode 100644
index 0000000..b2e1471
--- /dev/null
+++ b/src/prompt.rs
@@ -0,0 +1,310 @@
+use chrono;
+use regex;
+use std;
+
+use colors;
+use power;
+
+pub struct Prompt {
+    components_line1: Vec<Box<PromptComponent>>,
+    components_line2: Vec<Box<PromptComponent>>,
+    data: PromptData,
+    colors: colors::Colors,
+}
+
+trait PromptComponent {
+    fn display(&self, colors: &colors::Colors, data: &PromptData);
+    fn length(&self, data: &PromptData) -> usize;
+}
+
+struct PromptPath {
+    braces: String,
+}
+
+struct PromptBorder {
+    border: String,
+}
+
+struct PromptBattery {
+    length: usize,
+
+    braces: String,
+
+    full: String,
+    empty: String,
+    charging: String,
+    discharging: String,
+    unknown: String,
+}
+
+struct PromptIdentity {}
+
+struct PromptTime {
+    braces: String,
+}
+
+struct PromptCommandError {}
+
+struct PromptPrompt {
+    user: String,
+    root: String,
+}
+
+#[derive(Debug)]
+pub struct PromptData {
+    pub shell: colors::ShellType,
+    pub error_code: u8,
+    pub hostname: Option<String>,
+    pub terminal_cols: Option<usize>,
+    pub pwd: Option<std::path::PathBuf>,
+    pub home: Option<std::path::PathBuf>,
+    pub user: Option<String>,
+    pub is_root: bool,
+    pub time: chrono::DateTime<chrono::Local>,
+    pub power_info: power::PowerInfo,
+}
+
+impl Prompt {
+    pub fn new(data: PromptData) -> Prompt {
+        let shell = data.shell.clone();
+        let components_line1: Vec<Box<PromptComponent>> = vec![
+            Box::new(PromptPath { braces: "()".to_string() }),
+            Box::new(PromptBorder { border: "-".to_string() }),
+            Box::new(PromptBattery {
+                length: 10,
+                braces: "{}".to_string(),
+                full: "=".to_string(),
+                empty: "-".to_string(),
+                charging: "<".to_string(),
+                discharging: ">".to_string(),
+                unknown: "?".to_string(),
+            }),
+            Box::new(PromptIdentity {}),
+            Box::new(PromptTime { braces: "[]".to_string() }),
+        ];
+        let components_line2: Vec<Box<PromptComponent>> = vec![
+            Box::new(PromptCommandError {}),
+            Box::new(PromptPrompt {
+                user: "$".to_string(),
+                root: "#".to_string(),
+            }),
+        ];
+
+        Prompt {
+            components_line1: components_line1,
+            components_line2: components_line2,
+            data: data,
+            colors: colors::Colors::new(shell),
+        }
+    }
+
+    pub fn display(&self) {
+        self.colors.pad(1);
+        for component in self.components_line1.iter() {
+            component.display(&self.colors, &self.data);
+            self.colors.pad(1);
+        }
+
+        self.colors.newline();
+
+        for component in self.components_line2.iter() {
+            component.display(&self.colors, &self.data);
+            self.colors.pad(1);
+        }
+    }
+}
+
+impl PromptComponent for PromptPath {
+    fn display(&self, colors: &colors::Colors, data: &PromptData) {
+        colors.print_host(
+            &data.hostname,
+            &self.braces[0..1]
+        );
+        if let Some(ref pwd) = data.pwd {
+            colors.print(
+                "default",
+                &compress_path(pwd, &data.home, 20)
+            );
+        }
+        else {
+            colors.print("error", "???");
+        }
+        colors.print_host(
+            &data.hostname,
+            &self.braces[1..2]
+        );
+    }
+
+    fn length(&self, data: &PromptData) -> usize {
+        0
+    }
+}
+
+impl PromptComponent for PromptBorder {
+    fn display(&self, colors: &colors::Colors, data: &PromptData) {
+        colors.print("default", &self.border.repeat(20))
+    }
+
+    fn length(&self, data: &PromptData) -> usize {
+        0
+    }
+}
+
+impl PromptComponent for PromptBattery {
+    fn display(&self, colors: &colors::Colors, data: &PromptData) {
+        colors.print_host(
+            &data.hostname,
+            &self.braces[0..1]
+        );
+
+        if let Some(usage) = data.power_info.battery_usage() {
+            let color = if usage >= 0.8 {
+                "battery_full"
+            }
+            else if data.power_info.charging() {
+                "default"
+            }
+            else if usage >= 0.4 {
+                "default"
+            }
+            else if usage >= 0.15 {
+                "battery_warn"
+            }
+            else if usage >= 0.05 {
+                "battery_crit"
+            }
+            else {
+                "battery_emerg"
+            };
+            let filled = (self.length as f64 * usage).ceil() as usize;
+            let unfilled = self.length - filled;
+
+            if unfilled > 0 {
+                colors.print(
+                    color,
+                    &self.empty.repeat(unfilled)
+                );
+            }
+
+            if data.power_info.charging() {
+                colors.print(
+                    color,
+                    &self.charging
+                )
+            }
+            else {
+                colors.print(
+                    color,
+                    &self.discharging
+                )
+            }
+
+            if filled > 1 {
+                colors.print(
+                    color,
+                    &self.full.repeat(filled - 1)
+                )
+            }
+        }
+        else {
+            colors.print(
+                "battery_emerg",
+                &self.unknown.repeat(self.length)
+            )
+        }
+
+        colors.print_host(
+            &data.hostname,
+            &self.braces[1..2]
+        );
+    }
+
+    fn length(&self, data: &PromptData) -> usize {
+        0
+    }
+}
+
+impl PromptComponent for PromptIdentity {
+    fn display(&self, colors: &colors::Colors, data: &PromptData) {
+        colors.print_user(
+            &data.user,
+            &data.user.clone().unwrap_or("???".to_string())
+        );
+        colors.print("default", "@");
+        colors.print_host(
+            &data.hostname,
+            &data.hostname.clone().unwrap_or("???".to_string())
+        );
+    }
+
+    fn length(&self, data: &PromptData) -> usize {
+        0
+    }
+}
+
+impl PromptComponent for PromptTime {
+    fn display(&self, colors: &colors::Colors, data: &PromptData) {
+        colors.print_host(
+            &data.hostname,
+            &self.braces[0..1]
+        );
+
+        colors.print(
+            "default",
+            &format!("{}", data.time.format("%H:%M:%S"))
+        );
+
+        colors.print_host(
+            &data.hostname,
+            &self.braces[1..2]
+        );
+    }
+
+    fn length(&self, data: &PromptData) -> usize {
+        0
+    }
+}
+
+impl PromptComponent for PromptCommandError {
+    fn display(&self, colors: &colors::Colors, data: &PromptData) {
+        let color = if data.error_code == 0 {
+            "default"
+        }
+        else {
+            "error"
+        };
+        colors.print(color, &format!("{:03}", data.error_code));
+    }
+
+    fn length(&self, data: &PromptData) -> usize {
+        0
+    }
+}
+
+impl PromptComponent for PromptPrompt {
+    fn display(&self, colors: &colors::Colors, data: &PromptData) {
+        let prompt = if data.is_root {
+            &self.root
+        }
+        else {
+            &self.user
+        };
+        colors.print_user(&data.user, prompt);
+    }
+
+    fn length(&self, data: &PromptData) -> usize {
+        0
+    }
+}
+
+fn compress_path<T: AsRef<std::path::Path>, U: AsRef<std::path::Path>>(path: T, home: &Option<U>, max_len: u16) -> String {
+    let path_str = path.as_ref().to_string_lossy().into_owned();
+    if let &Some(ref home) = home {
+        let home_str = home.as_ref().to_string_lossy().into_owned();
+        let home_re = regex::Regex::new(&(r"^".to_string() + &regex::escape(&home_str))).unwrap();
+        home_re.replace(&path_str, "~").into_owned()
+    }
+    else {
+        path_str
+    }
+}
diff --git a/src/system_info.rs b/src/system_info.rs
new file mode 100644
index 0000000..2079ea7
--- /dev/null
+++ b/src/system_info.rs
@@ -0,0 +1,46 @@
+use chrono;
+use hostname;
+use term_size;
+use std;
+use users;
+
+use power;
+
+pub fn hostname() -> Option<String> {
+    hostname::get_hostname()
+}
+
+pub fn terminal_cols() -> Option<usize> {
+    if let Some((w, _h)) = term_size::dimensions() {
+        Some(w)
+    }
+    else {
+        None
+    }
+}
+
+pub fn pwd() -> Option<std::path::PathBuf> {
+    std::env::current_dir().ok()
+}
+
+pub fn home() -> Option<std::path::PathBuf> {
+    std::env::var("HOME")
+        .map(|dir| std::path::Path::new(&dir).to_path_buf())
+        .ok()
+}
+
+pub fn user() -> Option<String> {
+    users::get_current_username()
+}
+
+pub fn is_root() -> bool {
+    users::get_current_uid() == 0
+}
+
+pub fn time() -> chrono::DateTime<chrono::Local> {
+    chrono::Local::now()
+}
+
+pub fn power_info() -> power::PowerInfo {
+    power::PowerInfo::new()
+}