fix retain key

layercake/layercake/dynamodb.py

@@ -146,11 +146,15 @@ if TYPE_CHECKING:
             Optional specification for nested data extraction.
         remove_prefix: str, optional
             Optional prefix to remove from the key when forming the result dict.
+        retain_key: bool, optional
+            Use the key itself as value if True; otherwise, use the extracted value.
         """
 
         sk: str
         path_spec: str | None = None
         remove_prefix: str | None = None
+        retain_key: bool = False
+
 else:
 
     class SortKey(str):
@@ -166,6 +170,8 @@ else:
             Optional specification for nested data extraction.
         remove_prefix: str, optional
             Optional prefix to remove from the key when forming the result dict.
+        retain_key: bool, optional
+            Use the key itself as value if True; otherwise, use the extracted value.
         """
 
         def __new__(
@@ -174,6 +180,7 @@ else:
             *,
             path_spec: str | None = None,
             remove_prefix: str | None = None,
+            retain_key: bool = False,
         ) -> str:
             return super().__new__(cls, sk)
 
@@ -183,12 +190,14 @@ else:
             *,
             path_spec: str | None = None,
             remove_prefix: str | None = None,
+            retain_key: bool = False,
         ) -> None:
             # __init__ is used to store the parameters for later reference.
             # For immutable types like str, __init__ cannot change the instance's value.
             self.sk = sk
             self.path_spec = path_spec
             self.remove_prefix = remove_prefix
+            self.retain_key = retain_key
 
 
 class Key(ABC, dict):
@@ -929,29 +938,32 @@ class DynamoDBCollection:
         else:
             head, tail = {}, items
 
-        def _getin(pair: KeyPair, v: dict) -> dict:
-            v = omit((PK, SK), v)
+        def _getin(pair: KeyPair, obj: dict) -> dict:
+            obj = omit((PK, SK), obj)
             sk = pair[SK]
             path_spec = getattr(sk, 'path_spec', None)
 
             if path_spec:
                 from glom import glom
 
-                return glom(v, path_spec)
-            return v
+                return glom(obj, path_spec)
+            return obj
 
         def _removeprefix(pair: KeyPair) -> str:
+            pk = pair[PK]
             sk = pair[SK]
 
             if not isinstance(sk, SortKey):
-                return pair[PK]
+                return pk
 
-            return sk.removeprefix(sk.remove_prefix or '')
+            key = pk if sk.retain_key else sk
+
+            return key.removeprefix(sk.remove_prefix or '')
 
         return head | {
-            _removeprefix(pair): _getin(pair, item)
-            for pair, item in zip(sortkeys, tail)
-            if item
+            _removeprefix(pair): _getin(pair, obj)
+            for pair, obj in zip(sortkeys, tail)
+            if obj
         }
 
     def query(

layercake/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "layercake"
-version = "0.3.0"
+version = "0.3.1"
 description = "Packages shared dependencies to optimize deployment and ensure consistency across functions."
 readme = "README.md"
 authors = [

layercake/tests/test_dynamodb.py

@@ -3,7 +3,6 @@ from decimal import Decimal
 from ipaddress import IPv4Address
 
 import pytest
-
 from layercake.dateutils import ttl
 from layercake.dynamodb import (
     ComposeKey,
@@ -371,3 +370,23 @@ def test_collection_get_items_pair_unflatten(
         'cpf': {'user_id': '5OxmMjL-ujoR5IMGegQz'},
         'email': {'user_id': '5OxmMjL-ujoR5IMGegQz'},
     }
+
+
+def test_collection_get_items_pair_path_spec(
+    dynamodb_seeds,
+    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
+):
+    collect = DynamoDBCollection(dynamodb_persistence_layer)
+    doc = collect.get_items(
+        KeyPair('cpf', SortKey('07879819908', path_spec='user_id', retain_key=True))
+        + KeyPair(
+            'email',
+            SortKey('osergiosiqueira@gmail.com', path_spec='user_id', retain_key=True),
+        ),
+        flatten_top=False,
+    )
+
+    assert doc == {
+        'cpf': '5OxmMjL-ujoR5IMGegQz',
+        'email': '5OxmMjL-ujoR5IMGegQz',
+    }