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')

layercake/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "layercake"
-version = "0.1.1"
+version = "0.1.2"
 description = "Add your description here"
 readme = "README.md"
 authors = [

layercake/tests/test_dynamodb.py

@@ -171,18 +171,33 @@ def test_collection_get_items(
     collect = DynamoDBCollection(dynamodb_persistence_layer)
 
     # This data was added from seeds
-    data = collect.get_items(
+    logs = collect.get_items(
         PartitionKey(
             ComposeKey('5OxmMjL-ujoR5IMGegQz', prefix='logs'),
         ),
     )
-    assert len(data['items']) == 2
+    assert len(logs['items']) == 2
+    assert logs == {
+        'items': [
+            {
+                'sk': '2024-02-08T16:42:33.776409-03:00',
+                'action': 'OPEN_EMAIL',
+                'id': '5OxmMjL-ujoR5IMGegQz',
+            },
+            {
+                'sk': '2019-03-25T00:00:00-03:00',
+                'action': 'CLICK_EMAIL',
+                'id': '5OxmMjL-ujoR5IMGegQz',
+            },
+        ],
+        'last_key': None,
+    }
 
     # This data was added from seeds
-    data = collect.get_items(
+    emails = collect.get_items(
         KeyPair('5OxmMjL-ujoR5IMGegQz', PrefixKey('emails#')),
     )
-    assert data == {
+    assert emails == {
         'items': [
             {
                 'email_verified': True,

layercake/uv.lock

@@ -417,7 +417,7 @@ wheels = [
 
 [[package]]
 name = "layercake"
-version = "0.1.0"
+version = "0.1.1"
 source = { editable = "." }
 dependencies = [
     { name = "aws-lambda-powertools", extra = ["all"] },