3 files changed, 59 insertions, 5 deletions
diff --git a/README.md b/README.md
index ee3df48..7eebf3e 100644
--- a/README.md
+++ b/README.md
@@ -25,6 +25,12 @@ Save it to file:
 ```python
 from ymlstash import YmlStash
 
-stash = YmlStash("path/to/db")
-stash.save(user)
+stash = YmlStash(User, "path/to/db")
+stash.save("yuval", user)
+```
+
+Load from file:
+
+```python
+user = stash.load("yuval")
 ```
diff --git a/tests/test_stash.py b/tests/test_stash.py
index b87c882..defdd7e 100644
--- a/tests/test_stash.py
+++ b/tests/test_stash.py
@@ -1,9 +1,27 @@
 from pathlib import Path
 from ymlstash import YmlStash
+from dataclasses import dataclass
+
+
+@dataclass
+class User:
+    name: str
+    age: int
 
 
 def test_stash_path():
-    stash = YmlStash("foo")
+    stash = YmlStash(User, "foo")
     assert stash.path == Path("foo")
-    stash = YmlStash(Path("foo"))
+    stash = YmlStash(User, Path("foo"))
     assert stash.path == Path("foo")
+
+
+def test_stash():
+    stash = YmlStash(User, "/tmp/")
+    yuval = User(name="yuval", age=42)
+    stash.save("foo", yuval)
+    assert stash.list_all_keys() == ["foo"]
+    obj = stash.load("foo")
+    assert obj == yuval
+    stash.drop()
+    assert stash.list_all_keys() == []
diff --git a/ymlstash/stash.py b/ymlstash/stash.py
index f2839ef..d9c169b 100644
--- a/ymlstash/stash.py
+++ b/ymlstash/stash.py
@@ -1,6 +1,36 @@
+import yaml
+import os
+
+from dataclasses import asdict
 from pathlib import Path
 
 
 class YmlStash:
-    def __init__(self, path):
+    def __init__(self, clazz, path, file_suffix="yml", unsafe=False):
+        self.clazz = clazz
         self.path = Path(path)
+        self.file_suffix = f".{file_suffix}"
+        self.yaml_loader = yaml.SafeLoader
+
+    def _get_path(self, key):
+        return self.path / f"{key}{self.file_suffix}"
+
+    def load(self, key):
+        with open(self._get_path(key)) as f:
+            y = yaml.load(f.read(), Loader=self.yaml_loader)
+        return self.clazz(**y)
+
+    def save(self, key, obj):
+        with open(self._get_path(key), "w") as f:
+            f.write(yaml.dump(asdict(obj)))
+
+    def _list_all_files(self):
+        return [f for f in os.listdir(self.path) if f.endswith(self.file_suffix)]
+
+    def list_all_keys(self):
+        return [f.replace(self.file_suffix, "") for f in self._list_all_files()]
+
+    def drop(self):
+        for f in self._list_all_files():
+            if f.endswith(self.file_suffix):
+                os.remove(self.path / f)