From e6ee34846ed867a165197196c07a44b0587dfd84 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Niels=20Lerche=20S=C3=B8rensen?= <nielslerches@gmail.com>
Date: Tue, 20 Sep 2022 08:24:05 +0200
Subject: [PATCH] - Add primary key attributes to mapper

---
 datafiles/config.py            |  3 +++
 datafiles/mapper.py            | 12 ++++++++++++
 datafiles/tests/test_mapper.py |  1 +
 3 files changed, 16 insertions(+)

diff --git a/datafiles/config.py b/datafiles/config.py
index 5fb52e27..0ec8bdc6 100644
--- a/datafiles/config.py
+++ b/datafiles/config.py
@@ -15,6 +15,7 @@ class Meta:
     datafile_manual: bool = False
     datafile_defaults: bool = False
     datafile_infer: bool = False
+    datafile_key_attrs: Optional[Dict[str, Converter]] = None
 
 
 def load(obj) -> Meta:
@@ -30,5 +31,7 @@ def load(obj) -> Meta:
         meta.datafile_defaults = obj.Meta.datafile_defaults
     with suppress(AttributeError):
         meta.datafile_infer = obj.Meta.datafile_infer
+    with suppress(AttributeError):
+        meta.datafile_key_attrs = obj.Meta.datafile_key_attrs
 
     return meta
diff --git a/datafiles/mapper.py b/datafiles/mapper.py
index 26410964..e6f34d49 100644
--- a/datafiles/mapper.py
+++ b/datafiles/mapper.py
@@ -28,6 +28,7 @@ def __init__(
         manual: bool,
         defaults: bool,
         infer: bool,
+        key_attrs: Dict,
         root: Optional[Mapper] = None,
     ) -> None:
         assert manual is not None
@@ -38,6 +39,7 @@ def __init__(
         self._manual = manual
         self.defaults = defaults
         self._infer = infer
+        self.key_attrs = key_attrs
         self._last_load = 0.0
         self._last_data: Dict = {}
         self._root = root
@@ -278,6 +280,7 @@ def create_mapper(obj, root=None) -> Mapper:
     meta = config.load(obj)
     attrs = meta.datafile_attrs
     pattern = meta.datafile_pattern
+    key_attrs = meta.datafile_key_attrs
 
     if attrs is None and dataclasses.is_dataclass(obj):
         attrs = {}
@@ -287,6 +290,14 @@ def create_mapper(obj, root=None) -> Mapper:
             if pattern is None or self_name not in pattern:
                 attrs[field.name] = map_type(resolve(field.type, obj), name=field.name)  # type: ignore
 
+    if key_attrs is None and dataclasses.is_dataclass(obj):
+        key_attrs = {}
+        log.debug(f"Mapping key attributes for {obj.__class__} object")
+        for field in dataclasses.fields(obj):
+            self_name = f"self.{field.name}"
+            if pattern is not None and self_name in pattern:
+                key_attrs[field.name] = map_type(resolve(field.type, obj), name=field.name)  # type: ignore
+
     return Mapper(
         obj,
         attrs=attrs or {},
@@ -294,5 +305,6 @@ def create_mapper(obj, root=None) -> Mapper:
         manual=meta.datafile_manual,
         defaults=meta.datafile_defaults,
         infer=meta.datafile_infer,
+        key_attrs=key_attrs or {},
         root=root,
     )
diff --git a/datafiles/tests/test_mapper.py b/datafiles/tests/test_mapper.py
index be754aa5..0c062656 100644
--- a/datafiles/tests/test_mapper.py
+++ b/datafiles/tests/test_mapper.py
@@ -31,6 +31,7 @@ def mapper():
             manual=Meta.datafile_manual,
             defaults=Meta.datafile_defaults,
             infer=Meta.datafile_infer,
+            key_attrs={},
         )
 
     def describe_path():