- Add primary key attributes to mapper

Author: nielslerches

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

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,12 +290,21 @@ 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 {},
         pattern=pattern,
         manual=meta.datafile_manual,
         defaults=meta.datafile_defaults,
         infer=meta.datafile_infer,
+        key_attrs=key_attrs or {},
         root=root,
     )

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():