[Categories] User input added

Signed-off-by: Slava Barinov <rayslava@gmail.com>

Cargo.toml

@@ -14,4 +14,5 @@ csv = "1.1"
 structopt = "0.3"
 shellexpand = "2.1"
 radix_trie = { version = "0.2", features = ["serde"] }
+libc = "0.2"
 

src/categories.rs

@@ -1,3 +1,5 @@
+use crate::ui::input_category;
+use libc::isatty;
 use radix_trie::Trie;
 use serde::{Deserialize, Serialize};
 use std::cmp::Ordering;
@@ -31,6 +33,8 @@ impl PartialEq for CatStat {
 
 impl Eq for CatStat {}
 
+pub type CatStats = Trie<String, Vec<CatStat>>;
+
 /// Insert new `cat` into statistics vector or add single usage to existing cat
 fn update_stat(cat: &str, stat: &mut Vec<CatStat>) {
     let existing = stat.iter_mut().find(|stat| stat.category == cat);
@@ -47,7 +51,10 @@ fn update_stat(cat: &str, stat: &mut Vec<CatStat>) {
 
 /// Set up `cat` as category for `item`: update statistics or create new item
 /// in `storage`
-pub fn assign_category(item: &str, cat: &str, storage: &mut Trie<String, Vec<CatStat>>) {
+pub fn assign_category(item: &str, cat: &str, storage: &mut CatStats) {
+    if cat.is_empty() {
+        panic!("Do not assign empty category!")
+    }
     let existing = storage.get_mut(item);
     match existing {
         Some(stat) => update_stat(cat, stat),
@@ -62,16 +69,35 @@ pub fn assign_category(item: &str, cat: &str, storage: &mut Trie<String, Vec<Cat
 }
 
 /// Return most probable category for provided `item`
-pub fn get_top_category<'a>(
-    item: &str,
-    storage: &'a Trie<String, Vec<CatStat>>,
-) -> Option<&'a str> {
+pub fn get_top_category<'a>(item: &str, storage: &'a CatStats) -> Option<&'a str> {
     match storage.get(item) {
         Some(stats) => Some(&stats[0].category),
         None => None,
     }
 }
 
+/// Choose proper category or ask user
+pub fn get_category(item: &str, storage: &mut CatStats) -> String {
+    let istty = unsafe { isatty(libc::STDOUT_FILENO as i32) } != 0;
+    if istty {
+        let cat = input_category(item);
+        if cat.is_empty() {
+            match get_top_category(item, storage) {
+                Some(cat) => String::from(cat),
+                None => String::new(),
+            }
+        } else {
+            assign_category(&item, &cat, storage);
+            cat
+        }
+    } else {
+        match get_top_category(item, storage) {
+            Some(cat) => String::from(cat),
+            None => String::new(),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/main.rs

@@ -13,8 +13,12 @@ use std::time::Duration;
 use structopt::StructOpt;
 
 mod categories;
+use categories::get_category;
+use categories::CatStats;
+
 mod import;
 mod receipt;
+mod ui;
 
 fn read_receipt(f: &str) -> receipt::Receipt {
     let json = fs::read_to_string(f).expect("Can't read file");
@@ -37,10 +41,14 @@ mod tests {
     }
 }
 
-fn gen_splits(items: &[receipt::Item]) -> Vec<Split> {
+fn gen_splits(items: &[receipt::Item], cs: &mut CatStats) -> Vec<Split> {
     let mut result: Vec<Split> = Vec::new();
     for i in items.iter() {
-        let t = Split::new().memo(i.name.as_str()).amount(i.sum).build();
+        let t = Split::new()
+            .memo(i.name.as_str())
+            .amount(i.sum)
+            .category(&get_category(i.name.as_str(), cs))
+            .build();
 
         result.push(t);
     }
@@ -111,7 +119,7 @@ fn main() {
         ),
     };
 
-    let catmap: Trie<String, Vec<categories::CatStat>> = match db.get("catmap") {
+    let mut catmap: CatStats = match db.get("catmap") {
         Some(v) => v,
         None => Trie::new(),
     };
@@ -122,7 +130,7 @@ fn main() {
     }
 
     let receipt = read_receipt(&args.filename);
-    let splits = gen_splits(&receipt.items);
+    let splits = gen_splits(&receipt.items, &mut catmap);
     let acc = Account::new()
         .name("Wallet")
         .account_type(AccountType::Cash)

src/ui.rs

@@ -0,0 +1,9 @@
+use std::io::{stdin, stdout, Write};
+
+pub fn input_category(item: &str) -> String {
+    let mut x = String::with_capacity(64);
+    print!("Category for '{}'? > ", item);
+    let _ = stdout().flush();
+    stdin().read_line(&mut x).expect("Error reading input");
+    String::from(x.trim_end_matches('\n'))
+}