2 files changed, 14 insertions, 6 deletions
diff --git a/src/lib.rs b/src/lib.rs
index 04da2ca..9de5626 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -28,3 +28,15 @@ pub struct Mark {
 pub struct Primstav {
     pub marks: BTreeMap<String, Mark>,
 }
+
+impl Primstav {
+    pub fn new(marks: BTreeMap<String, Mark>) -> Primstav {
+        Primstav { marks }
+    }
+
+    /// Get an entry for the given date if one exists.
+    pub fn get(&self, date: &chrono::NaiveDate) -> Option<&Mark> {
+        let key = date.format("%d%m").to_string();
+        self.marks.get(&key)
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index fa18e1d..b93b19e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -19,7 +19,6 @@ use serde_yaml;
 use std::io::BufReader;
 use std::iter::Iterator;
 use std::fs::File;
-use std::string::ToString;
 
 mod locale;
 
@@ -66,9 +65,7 @@ fn print_entry(m: &Mark) {
 
 fn main() {
     let f = BufReader::new(File::open("data/primstav.yml").unwrap());
-    let p = Primstav {
-        marks: serde_yaml::from_reader(f).unwrap(),
-    };
+    let p = Primstav::new(serde_yaml::from_reader(f).unwrap());
 
     let mut args = std::env::args().skip(1);
     let local_date = match args.next() {
@@ -79,8 +76,7 @@ fn main() {
     let locale = locale::get();
     println!("{}", capitalize(&locale.format_date(&local_date)));
 
-    let key = local_date.format("%d%m").to_string();
-    match p.marks.get(&key) {
+    match p.get(&local_date) {
         Some(d) => print_entry(&d),
         None => ()
     }