3 files changed, 35 insertions, 14 deletions

diff --git a/README.md b/README.md
index 9cfc571..25737d2 100644
--- a/README.md
+++ b/README.md
@@ -2,6 +2,8 @@
 
 A simple ORM-like utility for operating on local YAML files via Python dataclasses.
 
+## Usage
+
 Define a dataclass:
 
 ```python
@@ -16,7 +18,7 @@ class User:
     key: ClassVar[str] = "name"
 ```
 
-Note the `key` field which is used to denote that `name` should be used as the primary key field.
+Note the special `key` field which is used to denote that `name` should be used as the primary key field.
 
 Instantiate a new object:
 
@@ -30,7 +32,7 @@ Save it to file:
 from ymlstash import YmlStash
 
 stash = YmlStash(User, "path/to/db")
-stash.save("yuval", user)
+stash.save(user)
 ```
 
 Load from file:
diff --git a/tests/test_stash.py b/tests/test_stash.py
index e7e573b..6a11db9 100644
--- a/tests/test_stash.py
+++ b/tests/test_stash.py
@@ -1,3 +1,5 @@
+import pytest
+
 from pathlib import Path
 from ymlstash import YmlStash
 from dataclasses import dataclass
@@ -20,7 +22,7 @@ def test_stash_path():
 def test_stash():
     stash = YmlStash(User, "/tmp/")
     yuval = User(name="yuval", age=42)
-    stash.save("foo", yuval)
+    stash.save(yuval, "foo")
     assert stash.list_all_keys() == ["foo"]
     obj = stash.load("foo")
     assert obj == yuval
@@ -28,15 +30,26 @@ def test_stash():
     assert stash.list_all_keys() == []
 
 
-@dataclass
-class Dog:
-    name: str
-    key: ClassVar[str] = "name"
+def test_key_field():
+    @dataclass
+    class Rat:
+        key: ClassVar[str] = "name"
+
+    with pytest.raises(Exception):
+        YmlStash(Rat, ".")
 
+    @dataclass
+    class Dog:
+        name: str
+        key: ClassVar[str] = "name"
 
-def test_auto_key():
     stash = YmlStash(Dog, "/tmp/")
     terra = Dog(name="terra")
-    stash.save(None, terra)
+    stash.save(terra)
     assert stash.list_all_keys() == ["terra"]
+
+    terra = Dog(name="terra")
+    stash.save(terra, key="dupe")  # override key
+    assert stash.list_all_keys() == ["terra", "dupe"]
+
     stash.drop()
diff --git a/ymlstash/stash.py b/ymlstash/stash.py
index e20c432..55bf016 100644
--- a/ymlstash/stash.py
+++ b/ymlstash/stash.py
@@ -1,17 +1,23 @@
 import yaml
 import os
 
-from dataclasses import asdict
+from dataclasses import asdict, fields
 from pathlib import Path
 
 
 class YmlStash:
-    def __init__(self, clazz, path, key_field=None, file_suffix="yml", unsafe=False):
+    def __init__(self, clazz, path, file_suffix="yml", unsafe=False):
         self.clazz = clazz
         self.path = Path(path)
-        self.key_field = key_field
         self.file_suffix = f".{file_suffix}"
         self.yaml_loader = yaml.SafeLoader
+        self._validate()
+
+    def _validate(self):
+        key_field = getattr(self.clazz, "key", None)
+        field_names = [f.name for f in fields(self.clazz)]
+        if key_field and key_field not in field_names:
+            raise Exception(f"Dataclass {self.clazz} has no key field '{key_field}'")
 
     def _get_path(self, key):
         return self.path / f"{key}{self.file_suffix}"
@@ -21,11 +27,11 @@ class YmlStash:
             y = yaml.load(f.read(), Loader=self.yaml_loader)
         return self.clazz(**y)
 
-    def save(self, key, obj):
+    def save(self, obj, key=None):
         if not key:
             key_field = self.clazz.key
             if not key_field:
-                raise Exception("no key")
+                raise Exception("Cannot save object without a key or key field")
             key = getattr(obj, key_field)
         with open(self._get_path(key), "w") as f:
             f.write(yaml.dump(asdict(obj)))