Initial autokey functionality

3 files changed, 26 insertions, 1 deletions

diff --git a/README.md b/README.md
index 7eebf3e..9cfc571 100644
--- a/README.md
+++ b/README.md
@@ -6,14 +6,18 @@ Define a dataclass:
 
 ```python
 from dataclasses import dataclass
+from typing import ClassVar
 
 @dataclass
 class User:
     name: str
     age: int
     active: bool
+    key: ClassVar[str] = "name"
 ```
 
+Note the `key` field which is used to denote that `name` should be used as the primary key field.
+
 Instantiate a new object:
 
 ```python
diff --git a/tests/test_stash.py b/tests/test_stash.py
index defdd7e..e7e573b 100644
--- a/tests/test_stash.py
+++ b/tests/test_stash.py
@@ -1,6 +1,7 @@
 from pathlib import Path
 from ymlstash import YmlStash
 from dataclasses import dataclass
+from typing import ClassVar
 
 
 @dataclass
@@ -25,3 +26,17 @@ def test_stash():
     assert obj == yuval
     stash.drop()
     assert stash.list_all_keys() == []
+
+
+@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)
+    assert stash.list_all_keys() == ["terra"]
+    stash.drop()
diff --git a/ymlstash/stash.py b/ymlstash/stash.py
index d9c169b..e20c432 100644
--- a/ymlstash/stash.py
+++ b/ymlstash/stash.py
@@ -6,9 +6,10 @@ from pathlib import Path
 
 
 class YmlStash:
-    def __init__(self, clazz, path, file_suffix="yml", unsafe=False):
+    def __init__(self, clazz, path, key_field=None, 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
 
@@ -21,6 +22,11 @@ class YmlStash:
         return self.clazz(**y)
 
     def save(self, key, obj):
+        if not key:
+            key_field = self.clazz.key
+            if not key_field:
+                raise Exception("no key")
+            key = getattr(obj, key_field)
         with open(self._get_path(key), "w") as f:
             f.write(yaml.dump(asdict(obj)))