remove prefix for partition key

layercake/layercake/dynamodb.py

@@ -669,10 +669,12 @@ class DynamoDBCollection:
         items = res['items']
         last_key = _startkey_b64encode(res['last_key']) if res['last_key'] else None
 
-        # Remove prefix from Sort Key if `key[SK]` is a PrefixKey
-        if isinstance(key.get(SK), PrefixKey):
-            prefix = key[SK].prefix
-            items = [item | {SK: item[SK].removeprefix(prefix)} for item in items]
+        match key.get(PK), key.get(SK):
+            case ComposeKey(), _:  # Remove prefix from Partition Key
+                prefix = key[PK].prefix + key[PK].delimiter
+                items = _remove_prefix(items, PK, prefix)
+            case _, PrefixKey():  # Remove prefix from Sort Key
+                items = _remove_prefix(items, SK, key[SK].prefix)
 
         return {
             'items': items,
@@ -680,6 +682,16 @@ class DynamoDBCollection:
         }
 
 
+def _remove_prefix(
+    items: list[dict[str, Any]],
+    /,
+    key: str,
+    prefix: str,
+) -> list[dict[str, Any]]:
+    """Remove the given prefix from the value associated with key in each item."""
+    return [x | {key: x[key].removeprefix(prefix)} for x in items]
+
+
 def _startkey_b64encode(obj: dict) -> str:
     s = json.dumps(obj)
     b = urlsafe_b64encode(s.encode('utf-8')).decode('utf-8')