From 55e6860c88b55d4b640a2b5a9a4fe6cd21b9bd16 Mon Sep 17 00:00:00 2001
From: Magnus Runesson <Magnus.Runesson@ai.se>
Date: Tue, 21 Jul 2026 10:05:22 +0200
Subject: [PATCH] Support bracket notation for array access in metadata keys

---
 tests/utils/test_metadata.py            | 14 ++++++++++++++
 training_data_packer/processor/clean.py |  6 ++++--
 training_data_packer/utils/metadata.py  |  1 +
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/tests/utils/test_metadata.py b/tests/utils/test_metadata.py
index 4648afe..b37ae2b 100644
--- a/tests/utils/test_metadata.py
+++ b/tests/utils/test_metadata.py
@@ -158,6 +158,20 @@ def test_get_shard_size_documents(self):
                 17,
                 42,
             ],
+            [
+                "array_square_bracket",
+                {"a": [{"b": 5}, {"b": 7}]},
+                "a[0].b",
+                17,
+                5,
+            ],
+            [
+                "array_glom",
+                {"a": [{"b": 5}, {"b": 7}]},
+                "a.0.b",
+                17,
+                5,
+            ],
         ]
     )
     def test_get_metadata_value(self, name, metadata, key, default_value, expected):
diff --git a/training_data_packer/processor/clean.py b/training_data_packer/processor/clean.py
index a80df7b..8364963 100644
--- a/training_data_packer/processor/clean.py
+++ b/training_data_packer/processor/clean.py
@@ -20,8 +20,10 @@ def __next__(self):
         src_doc = next(self._src_data)
         for field in self._mapper:
             try:
-                src_doc[field] = glom.glom(src_doc, self._mapper[field])
-                glom.delete(src_doc, self._mapper[field])
+                key = self._mapper[field]
+                key = key.replace("[", ".").replace("]", "")
+                src_doc[field] = glom.glom(src_doc, key)
+                glom.delete(src_doc, key)
             except glom.PathAccessError as e:
                 if field not in src_doc:
                     raise e
diff --git a/training_data_packer/utils/metadata.py b/training_data_packer/utils/metadata.py
index 05c929c..790d51b 100644
--- a/training_data_packer/utils/metadata.py
+++ b/training_data_packer/utils/metadata.py
@@ -22,6 +22,7 @@ def get_metadata_value(metadata: dict[str, Any], key: str, default: Any = None)
     :return: The value extracted from the metadata corresponding to the given
              key, or the default value if the key is not resolved.
     """
+    key = key.replace("[", ".").replace("]", "")
     return glom.glom(metadata, key, default=default)