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(