summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--README.md4
-rw-r--r--tests/test_stash.py15
-rw-r--r--ymlstash/stash.py8
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)))