From b6ae078bade3be1c07a6fe99f73f01cb0cff26f4 Mon Sep 17 00:00:00 2001
From: Magnus Runesson <Magnus.Runesson@ai.se>
Date: Thu, 23 Jul 2026 11:53:23 +0200
Subject: [PATCH] Use Metadata class everywhere.

Still todo is move functions operating on metadata into Metadata class as methods.
---
 .../integration/test_collect_metrics_tree.py  |  13 +-
 tests/integration/test_propella_merge.py      |   3 +-
 tests/mode/test_release.py                    | 155 ++++++-----
 tests/processor/sample/test_sampler.py        |  45 +--
 tests/processor/test_clean.py                 |  13 +-
 tests/processor/test_filters.py               |   2 +-
 tests/processor/test_parallel_merger.py       |  61 +++--
 tests/processor/test_propella.py              |  27 +-
 tests/test_propella_structure.py              |   7 +-
 tests/utils/test_metadata.py                  | 258 ++++++++++--------
 training_data_packer/merge.py                 |   2 +-
 training_data_packer/mode/release.py          |   6 +-
 training_data_packer/mode/sample.py           |   7 +-
 training_data_packer/processor/clean.py       |   4 +-
 training_data_packer/processor/filters.py     |   2 +-
 .../processor/parallel_merger.py              |  14 +-
 training_data_packer/processor/pii_masking.py |   6 +-
 training_data_packer/processor/propella.py    |   4 +-
 .../processor/sample/sampler.py               |   4 +-
 training_data_packer/propella_merge.py        |   9 +-
 training_data_packer/propella_structure.py    |   8 +-
 training_data_packer/utils/file.py            |   2 +-
 training_data_packer/utils/metadata.py        |  47 +++-
 training_data_packer/utils/misc.py            |   2 +-
 24 files changed, 400 insertions(+), 301 deletions(-)

diff --git a/tests/integration/test_collect_metrics_tree.py b/tests/integration/test_collect_metrics_tree.py
index 2e4772e..9416b68 100644
--- a/tests/integration/test_collect_metrics_tree.py
+++ b/tests/integration/test_collect_metrics_tree.py
@@ -8,6 +8,7 @@
 import yaml
 
 from training_data_packer.collect_metrics import main
+from training_data_packer.utils.metadata import Metadata
 
 
 class TestCollectMetricsTree(unittest.TestCase):
@@ -16,12 +17,14 @@ def test_collect_metrics_tree(self):
             tmp_path = Path(tmp_dir)
 
             # Create metadata.yaml
-            metadata = {
-                "release": {"default": {"pack": "tree"}, "part1": {}, "part2": {}},
-                "source": {"part1": {}, "part2": {}},
-            }
+            metadata = Metadata(
+                {
+                    "release": {"default": {"pack": "tree"}, "part1": {}, "part2": {}},
+                    "source": {"part1": {}, "part2": {}},
+                }
+            )
             with open(tmp_path / "metadata.yaml", "w") as f:
-                yaml.dump(metadata, f)
+                yaml.dump(metadata.data, f)
 
             # Create release_raw structure
             release_raw = tmp_path / "release-raw"
diff --git a/tests/integration/test_propella_merge.py b/tests/integration/test_propella_merge.py
index f1d70c9..c2c3d3d 100644
--- a/tests/integration/test_propella_merge.py
+++ b/tests/integration/test_propella_merge.py
@@ -6,6 +6,7 @@
 from training_data_packer import propella_merge
 from training_data_packer.propella_merge import process_file
 from training_data_packer.utils.file import GenericJsonlReader, JsonlZstWriter
+from training_data_packer.utils.metadata import Metadata
 from training_data_packer.utils.metrics import read_metrics_from_file
 
 
@@ -15,7 +16,7 @@ def test_process_skip_existing_file(self):
         Test that process_file skips files that already exist in the output directory.
         """
         with TemporaryDirectory() as test_dir:
-            metadata = {"id": "id", "suffix": ".jsonl.zst"}
+            metadata = Metadata({"id": "id", "suffix": ".jsonl.zst"})
             source_name = "test_file.jsonl.zst"
             propella_dir = Path(test_dir).joinpath("propella-4b")
             propella_dir.mkdir(parents=True)
diff --git a/tests/mode/test_release.py b/tests/mode/test_release.py
index 03a4af1..c8b809d 100644
--- a/tests/mode/test_release.py
+++ b/tests/mode/test_release.py
@@ -1,6 +1,7 @@
 import unittest
 
 from training_data_packer.mode.release import parallel_package_pipeline
+from training_data_packer.utils.metadata import Metadata
 
 
 class TestParallelPackagePipeline(unittest.TestCase):
@@ -21,21 +22,23 @@ def test_first(self):
                 "tgt_lang": "als_Latn",
             }
         ]
-        metadata = {
-            "_internal": {"parallel": True},
-            "text": "text",
-            "parallel": {
-                "input": "source",
-                "source": {
-                    "text": "source_text",
-                    "language": "src_lang",
-                },
-                "target": {
-                    "text": "target_text",
-                    "language": "tgt_lang",
+        metadata = Metadata(
+            {
+                "_internal": {"parallel": True},
+                "text": "text",
+                "parallel": {
+                    "input": "source",
+                    "source": {
+                        "text": "source_text",
+                        "language": "src_lang",
+                    },
+                    "target": {
+                        "text": "target_text",
+                        "language": "tgt_lang",
+                    },
                 },
-            },
-        }
+            }
+        )
         part_config = {}
         piis = iter([])
         contaminations = iter([])
@@ -71,21 +74,23 @@ def test_pii_without_id(self):
                 "tgt_lang": "als_Latn",
             },
         ]
-        metadata = {
-            "_internal": {"parallel": True},
-            "text": "text",
-            "parallel": {
-                "input": "source",
-                "source": {
-                    "text": "source_text",
-                    "language": "src_lang",
-                },
-                "target": {
-                    "text": "target_text",
-                    "language": "tgt_lang",
+        metadata = Metadata(
+            {
+                "_internal": {"parallel": True},
+                "text": "text",
+                "parallel": {
+                    "input": "source",
+                    "source": {
+                        "text": "source_text",
+                        "language": "src_lang",
+                    },
+                    "target": {
+                        "text": "target_text",
+                        "language": "tgt_lang",
+                    },
                 },
-            },
-        }
+            }
+        )
         part_config = {}
         piis = iter([{"hash": "7ea462d1386566606429b6ba9c3adfe39d742c742c835b464a8e8ccdfdf71452"}])
         contaminations = iter([])
@@ -122,22 +127,24 @@ def test_pii_with_id(self):
                 "tgt_lang": "als_Latn",
             },
         ]
-        metadata = {
-            "_internal": {"parallel": True},
-            "id": "id",
-            "text": "text",
-            "parallel": {
-                "input": "source",
-                "source": {
-                    "text": "source_text",
-                    "language": "src_lang",
-                },
-                "target": {
-                    "text": "target_text",
-                    "language": "tgt_lang",
+        metadata = Metadata(
+            {
+                "_internal": {"parallel": True},
+                "id": "id",
+                "text": "text",
+                "parallel": {
+                    "input": "source",
+                    "source": {
+                        "text": "source_text",
+                        "language": "src_lang",
+                    },
+                    "target": {
+                        "text": "target_text",
+                        "language": "tgt_lang",
+                    },
                 },
-            },
-        }
+            }
+        )
         part_config = {}
         piis = iter([{"id": "2"}])
         contaminations = iter([])
@@ -172,21 +179,23 @@ def test_contamination_without_id(self):
                 "tgt_lang": "als_Latn",
             },
         ]
-        metadata = {
-            "_internal": {"parallel": True},
-            "text": "text",
-            "parallel": {
-                "input": "source",
-                "source": {
-                    "text": "source_text",
-                    "language": "src_lang",
-                },
-                "target": {
-                    "text": "target_text",
-                    "language": "tgt_lang",
+        metadata = Metadata(
+            {
+                "_internal": {"parallel": True},
+                "text": "text",
+                "parallel": {
+                    "input": "source",
+                    "source": {
+                        "text": "source_text",
+                        "language": "src_lang",
+                    },
+                    "target": {
+                        "text": "target_text",
+                        "language": "tgt_lang",
+                    },
                 },
-            },
-        }
+            }
+        )
         part_config = {}
         piis = iter([])
         contaminations = iter([{"hash": "7ea462d1386566606429b6ba9c3adfe39d742c742c835b464a8e8ccdfdf71452"}])
@@ -223,22 +232,24 @@ def test_contamination_with_id(self):
                 "tgt_lang": "als_Latn",
             },
         ]
-        metadata = {
-            "_internal": {"parallel": True},
-            "id": "id",
-            "text": "text",
-            "parallel": {
-                "input": "source",
-                "source": {
-                    "text": "source_text",
-                    "language": "src_lang",
-                },
-                "target": {
-                    "text": "target_text",
-                    "language": "tgt_lang",
+        metadata = Metadata(
+            {
+                "_internal": {"parallel": True},
+                "id": "id",
+                "text": "text",
+                "parallel": {
+                    "input": "source",
+                    "source": {
+                        "text": "source_text",
+                        "language": "src_lang",
+                    },
+                    "target": {
+                        "text": "target_text",
+                        "language": "tgt_lang",
+                    },
                 },
-            },
-        }
+            }
+        )
         part_config = {}
         piis = iter([])
         contaminations = iter([{"id": "1"}])
diff --git a/tests/processor/sample/test_sampler.py b/tests/processor/sample/test_sampler.py
index 4c36eca..f4151a1 100644
--- a/tests/processor/sample/test_sampler.py
+++ b/tests/processor/sample/test_sampler.py
@@ -6,6 +6,7 @@
 from parameterized import parameterized
 
 from training_data_packer.processor.sample.sampler import DynamicSampler, convert_to_type, sampler_factory
+from training_data_packer.utils.metadata import Metadata
 
 
 class TestDynamicSampler(unittest.TestCase):
@@ -75,52 +76,54 @@ def test_convert_to_type_preserves_dict(self):
 
 class TestSamplerFactory(unittest.TestCase):
     def test_sample_not_set(self):
-        metadata = {"release": {"foo": {}}}
-        src_file = "blabla/foo/bla"
+        metadata = Metadata({"release": {"foo": {}}})
+        src_file = Path("blabla/foo/bla")
         in_data = iter([])
         result_iter, _ = sampler_factory(in_data, metadata, src_file)
         self.assertEqual(in_data, result_iter)
 
     def test_sample_is_full(self):
-        metadata = {"release": {"foo": {"sample": "full"}}}
-        src_file = "blabla/foo/bla"
+        metadata = Metadata({"release": {"foo": {"sample": "full"}}})
+        src_file = Path("blabla/foo/bla")
         in_data = iter([])
         result_iter, _ = sampler_factory(in_data, metadata, src_file)
         self.assertEqual(in_data, result_iter)
 
     def test_sample_is_random(self):
-        metadata = {"release": {"foo": {"sample": "random", "budget": "75%"}}}
-        src_file = "blabla/foo/bla"
+        metadata = Metadata({"release": {"foo": {"sample": "random", "budget": "75%"}}})
+        src_file = Path("blabla/foo/bla")
         in_data = iter([])
         result_iter, _ = sampler_factory(in_data, metadata, src_file)
         self.assertIsInstance(result_iter, itertools.filterfalse)
 
     def test_sample_is_wds_register(self):
-        metadata = {"release": {"foo": {"sample": "wds+register"}}}
-        src_file = "blabla/foo/bla"
+        metadata = Metadata({"release": {"foo": {"sample": "wds+register"}}})
+        src_file = Path("blabla/foo/bla")
         in_data = iter([])
         result_iter, _ = sampler_factory(in_data, metadata, src_file)
         self.assertIsInstance(result_iter, itertools.chain)
 
     def test_sample_is_dynamic(self):
-        metadata = {
-            "_internal": {"collection_dir": Path(os.getcwd())},
-            "release": {
-                "foo": {
-                    "sample": "dynamic",
-                    "parameters": {"para": "5"},
-                    "filter": "tests/resources/sampler/test_sampler.py",
-                }
-            },
-        }
-        src_file = "blabla/foo/bla"
+        metadata = Metadata(
+            {
+                "_internal": {"collection_dir": Path(os.getcwd())},
+                "release": {
+                    "foo": {
+                        "sample": "dynamic",
+                        "parameters": {"para": "5"},
+                        "filter": "tests/resources/sampler/test_sampler.py",
+                    }
+                },
+            }
+        )
+        src_file = Path("blabla/foo/bla")
         in_data = iter([])
         result_iter, _ = sampler_factory(in_data, metadata, src_file)
         self.assertIsInstance(result_iter, itertools.chain)
 
     def test_sample_illegall(self):
-        metadata = {"release": {"foo": {"sample": "illegal"}}}
-        src_file = "blabla/foo/bla"
+        metadata = Metadata({"release": {"foo": {"sample": "illegal"}}})
+        src_file = Path("blabla/foo/bla")
         in_data = iter([])
         with self.assertRaises(ValueError):
             sampler_factory(in_data, metadata, src_file)
diff --git a/tests/processor/test_clean.py b/tests/processor/test_clean.py
index 19a6843..f270bb4 100644
--- a/tests/processor/test_clean.py
+++ b/tests/processor/test_clean.py
@@ -3,6 +3,7 @@
 from parameterized import parameterized
 
 from training_data_packer.processor.clean import AlignFieldNames, field_scrubber_factory
+from training_data_packer.utils.metadata import Metadata
 
 
 class TestAlignFieldNames(unittest.TestCase):
@@ -11,7 +12,7 @@ def test_no_mapping_needed(self):
             {"id": "1234", "text": "Happy"},
             {"id": "1235", "text": "Gazonk"},
         ]
-        align_it = AlignFieldNames(iter(src_indata), {"id": "id", "text": "text"})
+        align_it = AlignFieldNames(iter(src_indata), Metadata({"id": "id", "text": "text"}))
         align_list = list(align_it)
         self.assertEqual(align_list, src_indata)
 
@@ -20,7 +21,7 @@ def test_no_mapping_fields(self):
             {"id": "1234", "text": "Happy"},
             {"id": "1235", "text": "Gazonk"},
         ]
-        align_it = AlignFieldNames(iter(src_indata), {})
+        align_it = AlignFieldNames(iter(src_indata), Metadata({}))
         align_list = list(align_it)
         self.assertEqual(align_list, src_indata)
 
@@ -33,7 +34,7 @@ def test_map_id(self):
             {"id": "1234", "text": "Happy"},
             {"id": "1235", "text": "Gazonk"},
         ]
-        align_it = AlignFieldNames(iter(src_indata), {"id": "warcid", "text": "text"})
+        align_it = AlignFieldNames(iter(src_indata), Metadata({"id": "warcid", "text": "text"}))
         align_list = list(align_it)
         self.assertEqual(align_list, expected)
 
@@ -46,7 +47,7 @@ def test_map_text(self):
             {"id": "1234", "text": "Happy"},
             {"id": "1235", "text": "Gazonk"},
         ]
-        align_it = AlignFieldNames(iter(src_indata), {"id": "id", "text": "context"})
+        align_it = AlignFieldNames(iter(src_indata), Metadata({"id": "id", "text": "context"}))
         align_list = list(align_it)
         self.assertEqual(align_list, expected)
 
@@ -59,7 +60,9 @@ def test_hierarchy_text(self):
             {"id": "1234", "metadata": {}, "text": "Happy"},
             {"id": "1235", "metadata": {}, "text": "Gazonk"},
         ]
-        align_it = AlignFieldNames(iter(src_indata), {"id": "metadata.int-id", "text": "context", "doc_s": "doc_score"})
+        align_it = AlignFieldNames(
+            iter(src_indata), Metadata({"id": "metadata.int-id", "text": "context", "doc_s": "doc_score"})
+        )
         align_list = list(align_it)
         self.assertEqual(align_list, expected)
 
diff --git a/tests/processor/test_filters.py b/tests/processor/test_filters.py
index d54e160..011aacc 100644
--- a/tests/processor/test_filters.py
+++ b/tests/processor/test_filters.py
@@ -5,7 +5,7 @@
 
 class TestFilterOnBlocklist(unittest.TestCase):
     def test_filtering_with_metrics(self):
-        blocklist = ["2", "4", "7"]
+        blocklist = {"2", "4", "7"}
         input_elements = [
             {"id": "1", "name": "block1"},
             {"id": "2", "name": "block1"},
diff --git a/tests/processor/test_parallel_merger.py b/tests/processor/test_parallel_merger.py
index 81a444d..812da05 100644
--- a/tests/processor/test_parallel_merger.py
+++ b/tests/processor/test_parallel_merger.py
@@ -1,16 +1,19 @@
 import unittest
 
 from training_data_packer.processor.parallel_merger import ParallelLanguageMerger, ParallelSyntheticId
+from training_data_packer.utils.metadata import Metadata
 
 
 class TestParallelLanguageMerger(unittest.TestCase):
     def setUp(self):
-        self.metadata = {
-            "parallel": {
-                "source": {"language": "src_lang", "text": "source_text"},
-                "target": {"language": "tgt_lang", "text": "target_text"},
+        self.metadata = Metadata(
+            {
+                "parallel": {
+                    "source": {"language": "src_lang", "text": "source_text"},
+                    "target": {"language": "tgt_lang", "text": "target_text"},
+                }
             }
-        }
+        )
         self.part_config = {"parallel": {"count": "40"}}
 
     def test_initialization_default_values(self):
@@ -30,7 +33,7 @@ def test_initialization_custom_metric_name(self):
         self.assertEqual(merger._metric_name, custom_metric)
 
     def test_initialization_with_missing_metadata_uses_defaults(self):
-        metadata = {}
+        metadata = Metadata({})
         merger = ParallelLanguageMerger(metadata, {})
         self.assertEqual(merger._src_lang, "src_lang")
         self.assertEqual(merger._source_text_col, "source_text")
@@ -94,12 +97,14 @@ def test_mapper_text_format(self):
         self.assertIn(expected_text, result["text"])
 
     def test_initialization_custom_metadata_fields(self):
-        custom_metadata = {
-            "parallel": {
-                "source": {"language": "source_language", "text": "src_text"},
-                "target": {"language": "target_language", "text": "tgt_text"},
+        custom_metadata = Metadata(
+            {
+                "parallel": {
+                    "source": {"language": "source_language", "text": "src_text"},
+                    "target": {"language": "target_language", "text": "tgt_text"},
+                }
             }
-        }
+        )
         custom_part_config = {"parallel": {"count": "3"}}
 
         mapper = ParallelLanguageMerger(custom_metadata, custom_part_config)
@@ -110,12 +115,14 @@ def test_initialization_custom_metadata_fields(self):
         self.assertEqual(mapper._documents_per_batch, 3)
 
     def test_mapper_with_custom_fields(self):
-        custom_metadata = {
-            "parallel": {
-                "source": {"language": "source_language", "text": "src_text"},
-                "target": {"language": "target_language", "text": "tgt_text"},
+        custom_metadata = Metadata(
+            {
+                "parallel": {
+                    "source": {"language": "source_language", "text": "src_text"},
+                    "target": {"language": "target_language", "text": "tgt_text"},
+                }
             }
-        }
+        )
 
         mapper = ParallelLanguageMerger(custom_metadata, self.part_config).get_mapper()
         docs = [{"source_language": "eng", "src_text": "Hello", "target_language": "fra", "tgt_text": "Bonjour"}]
@@ -158,23 +165,25 @@ def test_get_merge_iterator_empty_input(self):
 class TestParallelSyntheticId(unittest.TestCase):
     def test_create_synthetic_id_mapper_default_config(self):
         doc = {"source_text": "foo bar", "target_text": "gazonk"}
-        to_test = ParallelSyntheticId({})
+        to_test = ParallelSyntheticId(Metadata({}))
         result = to_test.get_mapper()(doc)
         self.assertEqual("ecefffc59192dd7f300f750650a3d21e7a2ea1c9854552ab2f83c7fb4986a08b", result["id"])
         self.assertEqual({"parallel_synthetic_id": {"processed_records": 1}}, to_test.get_metrics())
 
     def test_create_synthetic_id_mapper_alternativ_fields(self):
         doc = {"field1": "foo bar", "field2": "gazonk"}
-        metadata = {
-            "parallel": {
-                "source": {
-                    "text": "field1",
-                },
-                "target": {
-                    "text": "field2",
-                },
+        metadata = Metadata(
+            {
+                "parallel": {
+                    "source": {
+                        "text": "field1",
+                    },
+                    "target": {
+                        "text": "field2",
+                    },
+                }
             }
-        }
+        )
         to_test = ParallelSyntheticId(metadata)
         result = to_test.get_mapper()(doc)
         self.assertEqual("ecefffc59192dd7f300f750650a3d21e7a2ea1c9854552ab2f83c7fb4986a08b", result["id"])
diff --git a/tests/processor/test_propella.py b/tests/processor/test_propella.py
index 5c497b4..be17456 100644
--- a/tests/processor/test_propella.py
+++ b/tests/processor/test_propella.py
@@ -7,6 +7,7 @@
     SourceToPropellaMapper,
     propella_annotate_factory,
 )
+from training_data_packer.utils.metadata import Metadata
 
 
 class TestPropellaProcessor(unittest.TestCase):
@@ -15,7 +16,7 @@ def test_mapper(self):
             {"my_id": "foo"},
             {"my_id": "bar"},
         ]
-        metadata = {"id": "my_id", "text": "text"}
+        metadata = Metadata({"id": "my_id", "text": "text"})
         lookup_data = {
             "foo": {"id": "foo", "data": "foo_data"},
             "bar": None,  # Emulating not found.
@@ -38,11 +39,13 @@ def test_mapper_with_id_hash(self):
             {"id": "foo", "text": "this is my document"},
             {"id": "bar", "text": "this is my second document"},
         ]
-        metadata = {
-            "id": "id",
-            "text": "text",
-            "propella-4b": {"hash": "sha256"},
-        }
+        metadata = Metadata(
+            {
+                "id": "id",
+                "text": "text",
+                "propella-4b": {"hash": "sha256"},
+            }
+        )
         lookup_data = {
             "0422bdf3f65ea9ebda3004d0b4e392906c10f5591a32034eeb41f86b78919d4d": {
                 "id": "0422bdf3f65ea9ebda3004d0b4e392906c10f5591a32034eeb41f86b78919d4d",
@@ -75,11 +78,13 @@ def test_mapper_with_id_hash_and_length(self):
             {"id": "foo", "text": "this is my document"},
             {"id": "bar", "text": "this is my second document"},
         ]
-        metadata = {
-            "id": "id",
-            "text": "text",
-            "propella-4b": {"hash": "sha256-32", "hash-id": "hash_id"},
-        }
+        metadata = Metadata(
+            {
+                "id": "id",
+                "text": "text",
+                "propella-4b": {"hash": "sha256-32", "hash-id": "hash_id"},
+            }
+        )
         lookup_data = {
             "0422bdf3f65ea9ebda3004d0b4e39290": {
                 "id": "0422bdf3f65ea9ebda3004d0b4e39290",
diff --git a/tests/test_propella_structure.py b/tests/test_propella_structure.py
index de56221..4f381fb 100644
--- a/tests/test_propella_structure.py
+++ b/tests/test_propella_structure.py
@@ -4,6 +4,7 @@
 from parameterized import parameterized
 
 from training_data_packer import propella_structure
+from training_data_packer.utils.metadata import Metadata
 
 
 class TestComputeOutputFilename(unittest.TestCase):
@@ -13,21 +14,21 @@ class TestComputeOutputFilename(unittest.TestCase):
                 "basic",
                 "/foo/bar/propella-4b",
                 "shard01/000.jsonl.zst",
-                {"suffix": ".jsonl.zst"},
+                Metadata({"suffix": ".jsonl.zst"}),
                 "/foo/bar/propella-4b/shard01/000.jsonl.zst",
             ],
             [
                 "zstd",
                 "/foo/bar/propella-4b",
                 "shard01/000.jsonl.zstd",
-                {"suffix": ".jsonl.zstd", "annotations": {"propella-4b": {"suffix": ".jsonl.zst"}}},
+                Metadata({"suffix": ".jsonl.zstd", "annotations": {"propella-4b": {"suffix": ".jsonl.zst"}}}),
                 "/foo/bar/propella-4b/shard01/000.jsonl.zst",
             ],
             [
                 "no_propella_suffix_not_zst",
                 "/foo/bar/propella-4b",
                 "shard01/000.jsonl.gz",
-                {"suffix": ".jsonl.gz"},
+                Metadata({"suffix": ".jsonl.gz"}),
                 "/foo/bar/propella-4b/shard01/000.jsonl.zst",
             ],
         ]
diff --git a/tests/utils/test_metadata.py b/tests/utils/test_metadata.py
index 4648afe..c8d4b12 100644
--- a/tests/utils/test_metadata.py
+++ b/tests/utils/test_metadata.py
@@ -5,6 +5,7 @@
 
 import training_data_packer
 from training_data_packer.utils.metadata import (
+    Metadata,
     calculate_file_path,
     get_matching_part,
     get_metadata_value,
@@ -14,125 +15,146 @@
 
 class TestMetadata(unittest.TestCase):
     def test_metadata_without_defaults(self):
-        indata = {"release": {"foo": {"sample": "full"}, "bar": {"sample": "full"}}}
-        self.assertEqual(get_matching_part(indata, "bla/foo/shard01"), ({"sample": "full"}, "foo"))
+        indata = Metadata({"release": {"foo": {"sample": "full"}, "bar": {"sample": "full"}}})
+        self.assertEqual(get_matching_part(indata, Path("bla/foo/shard01")), ({"sample": "full"}, "foo"))
 
     def test_metadata_merge_defaults(self):
-        indata = {"release": {"default": {"pack": "tree"}, "foo": {"sample": "full"}, "bar": {"sample": "full"}}}
-        self.assertEqual(get_matching_part(indata, "bla/foo/shard01"), ({"sample": "full", "pack": "tree"}, "foo"))
+        indata = Metadata(
+            {"release": {"default": {"pack": "tree"}, "foo": {"sample": "full"}, "bar": {"sample": "full"}}}
+        )
+        self.assertEqual(
+            get_matching_part(indata, Path("bla/foo/shard01")), ({"sample": "full", "pack": "tree"}, "foo")
+        )
 
     def test_metadata_use_defaults(self):
-        indata = {
-            "release": {
-                "default": {
-                    "pack": "tree",
-                    "sample": "wds",
-                },
-                "foo": {},
-                "bar": {
-                    "sample": "full",
-                },
+        indata = Metadata(
+            {
+                "release": {
+                    "default": {
+                        "pack": "tree",
+                        "sample": "wds",
+                    },
+                    "foo": {},
+                    "bar": {
+                        "sample": "full",
+                    },
+                }
             }
-        }
-        self.assertEqual(get_matching_part(indata, "bla/foo/shard01"), ({"sample": "wds", "pack": "tree"}, "foo"))
+        )
+        self.assertEqual(get_matching_part(indata, Path("bla/foo/shard01")), ({"sample": "wds", "pack": "tree"}, "foo"))
 
     def test_metadata_use_defaults_look_in_hierarchy(self):
-        indata = {
-            "source": {
-                "foo": {"value": "do not use"},
-            },
-            "release": {
-                "default": {
-                    "input": "source",
-                    "pack": "tree",
-                    "sample": "wds",
+        indata = Metadata(
+            {
+                "source": {
+                    "foo": {"value": "do not use"},
                 },
-                "bar": {
-                    "sample": "full",
+                "release": {
+                    "default": {
+                        "input": "source",
+                        "pack": "tree",
+                        "sample": "wds",
+                    },
+                    "bar": {
+                        "sample": "full",
+                    },
                 },
-            },
-        }
+            }
+        )
         self.assertEqual(
-            get_matching_part(indata, "bla/foo/shard01"), ({"input": "source", "sample": "wds", "pack": "tree"}, "foo")
+            get_matching_part(indata, Path("bla/foo/shard01")),
+            ({"input": "source", "sample": "wds", "pack": "tree"}, "foo"),
         )
 
     def test_metadata_fail_find_in_hierarchy(self):
-        indata = {
-            "source": {},
-            "release": {
-                "default": {
-                    "input": "source",
-                    "pack": "tree",
-                    "sample": "wds",
-                },
-                "bar": {
-                    "sample": "full",
+        indata = Metadata(
+            {
+                "source": {},
+                "release": {
+                    "default": {
+                        "input": "source",
+                        "pack": "tree",
+                        "sample": "wds",
+                    },
+                    "bar": {
+                        "sample": "full",
+                    },
                 },
-            },
-        }
-        self.assertEqual((None, None), get_matching_part(indata, "bla/foo/shard01"))
+            }
+        )
+        self.assertEqual((None, None), get_matching_part(indata, Path("bla/foo/shard01")))
 
     def test_metadata_override_defaults(self):
-        indata = {
-            "release": {
-                "default": {
-                    "pack": "tree",
-                    "sample": "wds",
-                },
-                "foo": {
-                    "sample": "full",
-                },
-                "bar": {
-                    "sample": "full",
-                },
+        indata = Metadata(
+            {
+                "release": {
+                    "default": {
+                        "pack": "tree",
+                        "sample": "wds",
+                    },
+                    "foo": {
+                        "sample": "full",
+                    },
+                    "bar": {
+                        "sample": "full",
+                    },
+                }
             }
-        }
-        self.assertEqual(get_matching_part(indata, "bla/foo/shard01"), ({"sample": "full", "pack": "tree"}, "foo"))
+        )
+        self.assertEqual(
+            get_matching_part(indata, Path("bla/foo/shard01")), ({"sample": "full", "pack": "tree"}, "foo")
+        )
 
     def test_no_matching_part(self):
-        indata = {
-            "source": {
-                "default": {
-                    "pack": "tree",
-                    "sample": "wds",
-                },
-                "bar": {
-                    "sample": "full",
-                },
+        indata = Metadata(
+            {
+                "source": {
+                    "default": {
+                        "pack": "tree",
+                        "sample": "wds",
+                    },
+                    "bar": {
+                        "sample": "full",
+                    },
+                }
             }
-        }
-        self.assertEqual((None, None), get_matching_part(indata, "bla/foo/shard01", "source"))
+        )
+        self.assertEqual((None, None), get_matching_part(indata, Path("bla/foo/shard01"), "source"))
 
     @parameterized.expand(
         [
             [
                 "flat",
-                {
-                    "_internal": {"mode": "release"},
-                    "release": {
-                        "default": {
-                            "unimportant": 5,
+                Metadata(
+                    {
+                        "_internal": {"mode": "release"},
+                        "release": {
+                            "default": {
+                                "unimportant": 5,
+                            },
+                            "part2": {},
+                            "part1": {},
                         },
-                        "part2": {},
-                        "part1": {},
-                    },
-                },
+                    }
+                ),
                 ["part1", "part2"],
             ],
             [
                 "hierarchy",
-                {
-                    "_internal": {"mode": "release"},
-                    "release": {
-                        "default": {"unimportant": 5, "input": "another"},
-                        "part2": {},
-                        "part1": {},
-                    },
-                    "another": {
-                        "part3": {},
-                        "part1": {},
-                    },
-                },
+                Metadata(
+                    {
+                        "_internal": {"mode": "release"},
+                        "release": {
+                            "default": {"unimportant": 5, "input": "another"},
+                            "part2": {},
+                            "part1": {},
+                        },
+                        "another": {
+                            "part3": {},
+                            "part1": {},
+                        },
+                    }
+                ),
                 ["part1", "part2", "part3"],
             ],
         ]
@@ -149,15 +171,29 @@ def test_get_shard_size_documents(self):
 
     @parameterized.expand(
         [
-            ["Value_set", {"foo": "bar"}, "foo", "gazsonk", "bar"],
+            ["Value_set", Metadata({"foo": "bar"}), "foo", "gazsonk", "bar"],
             ["default", {"foo": "bar"}, "key", "gazonk", "gazonk"],
             [
                 "hierarchy",
-                {"foo": {"bar": 42}},
+                Metadata({"foo": {"bar": 42}}),
                 "foo.bar",
                 17,
                 42,
             ],
+            [
+                "array_square_bracket",
+                Metadata({"a": [{"b": 5}, {"b": 7}]}),
+                "a[0].b",
+                17,
+                5,
+            ],
+            [
+                "array_glom",
+                Metadata({"a": [{"b": 5}, {"b": 7}]}),
+                "a.0.b",
+                17,
+                5,
+            ],
         ]
     )
     def test_get_metadata_value(self, name, metadata, key, default_value, expected):
@@ -178,11 +214,13 @@ class TestCalculateFilePath(unittest.TestCase):
 
     def test_calculate_file_paths_default(self):
         src_file = self.src_dir.joinpath("shard01/file01.jsonl.zst")
-        metadata = {
-            "_internal": {"collection_dir": Path("/srv/data/gazonk"), "mode": "release"},
-            "release": {"default": {"input": "source"}},
-            "suffix": ".jsonl.zstd",
-        }
+        metadata = Metadata(
+            {
+                "_internal": {"collection_dir": Path("/srv/data/gazonk"), "mode": "release"},
+                "release": {"default": {"input": "source"}},
+                "suffix": ".jsonl.zstd",
+            }
+        )
         contamination_file = calculate_file_path(src_file, metadata, "release", self.contamination_dir)
         pii_file = calculate_file_path(src_file, metadata, "release", self.pii_dir)
         propella_file = calculate_file_path(src_file, metadata, "release", self.propella_dir)
@@ -195,12 +233,14 @@ def test_calculate_file_paths_default(self):
 
     def test_calculate_file_paths_source_suffix(self):
         src_file = self.src_dir.joinpath("shard01/file01.jsonl.gz")
-        metadata = {
-            "_internal": {"collection_dir": Path("/srv/data/gazonk"), "mode": "release"},
-            "release": {"default": {"input": "source"}},
-            "source": {"default": {"suffix": ".jsonl.gz"}},
-            "suffix": ".jsonl.zst",
-        }
+        metadata = Metadata(
+            {
+                "_internal": {"collection_dir": Path("/srv/data/gazonk"), "mode": "release"},
+                "release": {"default": {"input": "source"}},
+                "source": {"default": {"suffix": ".jsonl.gz"}},
+                "suffix": ".jsonl.zst",
+            }
+        )
 
         contamination_file = calculate_file_path(src_file, metadata, "release", self.contamination_dir)
         pii_file = calculate_file_path(src_file, metadata, "release", self.pii_dir)
@@ -215,12 +255,14 @@ def test_calculate_file_paths_source_suffix(self):
     def test_calculate_file_paths_sample_source_suffix(self):
         src_file = self.src_dir.joinpath("shard01/file01.jsonl.gz")
         output_dir = Path("/srv/data/gazonk/sample")
-        metadata = {
-            "_internal": {"collection_dir": Path("/srv/data/gazonk"), "mode": "sample"},
-            "sample": {"default": {"input": "source"}},
-            "source": {"default": {"suffix": ".jsonl.gz"}},
-            "suffix": ".jsonl.zst",
-        }
+        metadata = Metadata(
+            {
+                "_internal": {"collection_dir": Path("/srv/data/gazonk"), "mode": "sample"},
+                "sample": {"default": {"input": "source"}},
+                "source": {"default": {"suffix": ".jsonl.gz"}},
+                "suffix": ".jsonl.zst",
+            }
+        )
 
         contamination_file = calculate_file_path(src_file, metadata, "sample", self.contamination_dir)
         pii_file = calculate_file_path(src_file, metadata, "sample", self.pii_dir)
diff --git a/training_data_packer/merge.py b/training_data_packer/merge.py
index 03c8669..90305eb 100644
--- a/training_data_packer/merge.py
+++ b/training_data_packer/merge.py
@@ -115,7 +115,7 @@ def process(collection_dir: Path, part: None | str = None, workers: int = 1, slu
         if parts == ["default"]:
             metadata["suffix"] = ".jsonl.zst"
             files = find_files(input_dir, metadata["suffix"])
-            docs_per_shard = get_shard_size_documents(get_metadata_value(metadata, "release.default"))
+            docs_per_shard = get_shard_size_documents(metadata["release.default"])
             merge(
                 files,
                 output_dir,
diff --git a/training_data_packer/mode/release.py b/training_data_packer/mode/release.py
index 1b7c37f..1c8ecae 100644
--- a/training_data_packer/mode/release.py
+++ b/training_data_packer/mode/release.py
@@ -12,12 +12,12 @@
 from training_data_packer.processor.sample.sampler import sampler_factory
 from training_data_packer.utils import metrics
 from training_data_packer.utils.file import GenericJsonlReader, JsonlZstWriter
-from training_data_packer.utils.metadata import get_matching_part, get_metadata_value
+from training_data_packer.utils.metadata import Metadata, get_matching_part, get_metadata_value
 
 
 def parallel_package_pipeline(
     src_iter: Iterable[dict[str, Any]],
-    metadata: dict[str, Any],
+    metadata: Metadata,
     part_config: dict[str, Any],
     piis: Iterable[dict[str, Any]],
     contaminations: Iterable[dict[str, Any]],
@@ -47,7 +47,7 @@ def parallel_package_pipeline(
 
 
 def package_file(
-    src_file: Path, metadata: dict, contamination_file: Path, pii_file: Path, propella_file: Path, out_file: Path
+    src_file: Path, metadata: Metadata, contamination_file: Path, pii_file: Path, propella_file: Path, out_file: Path
 ) -> None:
     tmp_out_file = out_file.parent.joinpath("." + out_file.name)
     if out_file.exists():
diff --git a/training_data_packer/mode/sample.py b/training_data_packer/mode/sample.py
index e316b1b..1b4cb91 100644
--- a/training_data_packer/mode/sample.py
+++ b/training_data_packer/mode/sample.py
@@ -1,5 +1,4 @@
 import os
-from argparse import ArgumentError
 from pathlib import Path
 
 from loguru import logger
@@ -9,10 +8,10 @@
 from training_data_packer.processor.sample.sampler import sampler_factory
 from training_data_packer.utils import metrics
 from training_data_packer.utils.file import GenericJsonlReader, JsonlZstWriter
-from training_data_packer.utils.metadata import get_matching_part
+from training_data_packer.utils.metadata import Metadata, get_matching_part
 
 
-def sample_file(src_file: Path, metadata: dict, propella_file: Path, out_file: Path) -> None:
+def sample_file(src_file: Path, metadata: Metadata, propella_file: Path, out_file: Path) -> None:
     tmp_out_file = out_file.parent.joinpath("." + out_file.name)
     if out_file.exists():
         # File is already processed. Do not process it again
@@ -40,7 +39,7 @@ def sample_file(src_file: Path, metadata: dict, propella_file: Path, out_file: P
     if "sample" in metadata:
         sampled_iter, sampler_metrics = sampler_factory(propella_iter, metadata, src_file, section_name="sample")
     else:
-        raise ArgumentError("Metadata does not contain sample section.")
+        raise ValueError("Metadata does not contain sample section.")
 
     writer = JsonlZstWriter(tmp_out_file)
     writer.write(sampled_iter)
diff --git a/training_data_packer/processor/clean.py b/training_data_packer/processor/clean.py
index a80df7b..7cfcd2a 100644
--- a/training_data_packer/processor/clean.py
+++ b/training_data_packer/processor/clean.py
@@ -3,9 +3,11 @@
 
 import glom
 
+from training_data_packer.utils.metadata import Metadata
+
 
 class AlignFieldNames:
-    def __init__(self, src_data: Iterator[Any], metadata: dict):
+    def __init__(self, src_data: Iterator[Any], metadata: Metadata):
         self._src_data = src_data
         self._mapper = {}
         if "id" in metadata and metadata["id"] != "id":
diff --git a/training_data_packer/processor/filters.py b/training_data_packer/processor/filters.py
index ee0cae6..fb122f4 100644
--- a/training_data_packer/processor/filters.py
+++ b/training_data_packer/processor/filters.py
@@ -8,7 +8,7 @@ class FilterOnBlocklist:
     Filter for blocklist. Filters out documents with ids in blocklist.
     """
 
-    def __init__(self, name: str, block_list: list[Any]):
+    def __init__(self, name: str, block_list: set[Any]):
         self.counter = 0
         self.name = name
         self.block_list = {str(x) for x in block_list}
diff --git a/training_data_packer/processor/parallel_merger.py b/training_data_packer/processor/parallel_merger.py
index 53de01a..595e868 100644
--- a/training_data_packer/processor/parallel_merger.py
+++ b/training_data_packer/processor/parallel_merger.py
@@ -3,7 +3,9 @@
 from collections.abc import Callable, Iterable
 from typing import Any
 
-from training_data_packer.utils.metadata import get_metadata_value
+import glom
+
+from training_data_packer.utils.metadata import Metadata, get_metadata_value
 from training_data_packer.utils.misc import hash_factory, lang_to_name
 
 
@@ -21,7 +23,7 @@ class ParallelLanguageMerger:
 
     def __init__(
         self,
-        metadata: dict[str, Any],
+        metadata: Metadata,
         part_config: dict[str, Any],
         flip_fn: callable[None, bool] = lambda: random.random() < 0.5,
         metric_name: str = "parallel_merger_matching",
@@ -35,7 +37,7 @@ def __init__(
         self._source_text_col = get_metadata_value(metadata, "parallel.source.text", "source_text")
         self._tgt_lang = get_metadata_value(metadata, "parallel.target.language", "tgt_lang")
         self._target_text_col = get_metadata_value(metadata, "parallel.target.text", "target_text")
-        self._documents_per_batch = int(get_metadata_value(part_config, "parallel.count", "40"))
+        self._documents_per_batch = int(glom.glom(part_config, "parallel.count", default="40"))
 
     def get_metrics(self):
         """
@@ -97,14 +99,14 @@ class ParallelSyntheticId:
 
     def __init__(
         self,
-        metadata: dict[str, Any],
+        metadata: Metadata,
         metric_name: str = "parallel_synthetic_id",
     ):
         self._metric_name = metric_name
         self._processed_records = 0
         hash_fn = hash_factory("sha256")
-        source_text_col = get_metadata_value(metadata, "parallel.source.text", "source_text")
-        target_text_col = get_metadata_value(metadata, "parallel.target.text", "target_text")
+        source_text_col = metadata.get("parallel.source.text", "source_text")
+        target_text_col = metadata.get("parallel.target.text", "target_text")
         self._id_fn = lambda doc: hash_fn(f"{doc[source_text_col]}{doc[target_text_col]}")
 
     def get_metrics(self):
diff --git a/training_data_packer/processor/pii_masking.py b/training_data_packer/processor/pii_masking.py
index fb5d1a6..e599bda 100644
--- a/training_data_packer/processor/pii_masking.py
+++ b/training_data_packer/processor/pii_masking.py
@@ -245,7 +245,7 @@ def _mask_ip_address(document: dict[str, Any], pii_record: dict[str, Any]) -> di
 
 
 def multilingual_mask_document(
-    document: dict[str, Any], pii_records: list[dict[str, Any]], mask_fields: list[str] = None
+    document: dict[str, Any], pii_records: list[dict[str, Any]], mask_fields: None | list[str] = None
 ) -> dict[str, Any]:
     """
     Masking a complete document. Masking pii records in document by replacing them with scrambled values.
@@ -288,7 +288,7 @@ def multilingual_mask_document(
 
 
 def openai_mask_document(
-    document: dict[str, Any], pii_records: list[dict[str, Any]], mask_fields: list[str] = None
+    document: dict[str, Any], pii_records: list[dict[str, Any]], mask_fields: None | list[str] = None
 ) -> dict[str, Any]:
     """
     Masking a complete document. Masking pii records are produced py OpenAIPrivacy tool and wrapper.
@@ -347,7 +347,7 @@ class PIIMasker:
     def __init__(
         self,
         masker_fn=multilingual_mask_document,
-        part_config: dict[str, Any] = None,
+        part_config: None | dict[str, Any] = None,
         metric_name: str = "pii_masker",
     ) -> None:
         self._metric_name = metric_name
diff --git a/training_data_packer/processor/propella.py b/training_data_packer/processor/propella.py
index e4664b0..0dbfdb1 100644
--- a/training_data_packer/processor/propella.py
+++ b/training_data_packer/processor/propella.py
@@ -4,7 +4,7 @@
 from loguru import logger
 
 import training_data_packer.utils.misc
-from training_data_packer.utils.metadata import get_metadata_value
+from training_data_packer.utils.metadata import Metadata, get_metadata_value
 
 
 class SourceToPropellaMapper:
@@ -22,7 +22,7 @@ class SourceToPropellaMapper:
 
     def __init__(
         self,
-        metadata: dict[str, Any],
+        metadata: Metadata,
         lookup_fn: Callable[[Any], dict[str, Any]],
         metric_name: str = "propella_matching",
     ):
diff --git a/training_data_packer/processor/sample/sampler.py b/training_data_packer/processor/sample/sampler.py
index 7efa5d6..267c284 100644
--- a/training_data_packer/processor/sample/sampler.py
+++ b/training_data_packer/processor/sample/sampler.py
@@ -9,7 +9,7 @@
 from loguru import logger
 
 from training_data_packer.processor.sample import sample_register
-from training_data_packer.utils.metadata import get_matching_part, get_metadata_value
+from training_data_packer.utils.metadata import Metadata, get_matching_part, get_metadata_value
 
 
 def read_sampler_fn(filename: Path | None):
@@ -113,7 +113,7 @@ def mapper(record: dict[str, Any]) -> list[dict[str, Any]]:
 
 
 def sampler_factory(
-    data_iterator: Iterable[dict[str, Any]], metadata: dict, src_file_name: Path, section_name: str = "release"
+    data_iterator: Iterable[dict[str, Any]], metadata: Metadata, src_file_name: Path, section_name: str = "release"
 ) -> tuple[Iterable[dict[str, Any]], Any]:
     part_config, part_name = get_matching_part(metadata, src_file_name, section_name)
     if part_config is None or "sample" not in part_config:
diff --git a/training_data_packer/propella_merge.py b/training_data_packer/propella_merge.py
index a71dc7b..4d1a8e0 100644
--- a/training_data_packer/propella_merge.py
+++ b/training_data_packer/propella_merge.py
@@ -2,7 +2,6 @@
 import os
 from concurrent.futures import ProcessPoolExecutor
 from pathlib import Path
-from typing import Any
 
 from loguru import logger
 
@@ -15,7 +14,7 @@
     find_files,
     get_subdirectories,
 )
-from training_data_packer.utils.metadata import get_metadata_value, read_metadata
+from training_data_packer.utils.metadata import Metadata, get_metadata_value, read_metadata
 from training_data_packer.utils.metrics import read_metrics_from_file
 from training_data_packer.utils.slurm import get_my_slurm_tasks
 
@@ -51,7 +50,7 @@ def process(collection_dir: Path, part: str = "", workers=1, slurm: bool = False
             process_file(metadata, src_name, output_dir)
 
 
-def process_file(metadata: dict[str, Any], source_name: str, propella_dir: Path):
+def process_file(metadata: Metadata, source_name: str, propella_dir: Path):
     """
     Processes and merges JSONL files from subdirectories within the specified
     propella directory. The function constructs the output path and checks for
@@ -73,8 +72,8 @@ def process_file(metadata: dict[str, Any], source_name: str, propella_dir: Path)
     """
     new_name = change_suffix(
         source_name,
-        get_metadata_value(metadata, "source.default.suffix", metadata["suffix"]),
-        get_metadata_value(metadata, "annotations.propella-4b.suffix", metadata["suffix"]),
+        metadata.get("source.default.suffix", metadata["suffix"]),
+        metadata.get("annotations.propella-4b.suffix", metadata["suffix"]),
     )
     out_file_name = propella_dir.joinpath(new_name)
 
diff --git a/training_data_packer/propella_structure.py b/training_data_packer/propella_structure.py
index 66c09bb..e5de80a 100644
--- a/training_data_packer/propella_structure.py
+++ b/training_data_packer/propella_structure.py
@@ -10,7 +10,7 @@
 from training_data_packer.storage.propella import get_lookup_fn
 from training_data_packer.utils import metrics
 from training_data_packer.utils.file import GenericJsonlReader, JsonlZstWriter, change_suffix, find_files
-from training_data_packer.utils.metadata import get_metadata_value, read_metadata
+from training_data_packer.utils.metadata import Metadata, get_metadata_value, read_metadata
 from training_data_packer.utils.slurm import get_my_slurm_tasks
 
 
@@ -43,17 +43,17 @@ def process(collection_dir: Path, propella_dir: Path, part: str = "", slurm: boo
             process_file(metadata, propella_lookup_fn, source_file, output_file, propella_metrics)
 
 
-def _compute_output_filename(output_dir: Path, rel_path: Path, metadata: dict[str, Any]) -> Path:
+def _compute_output_filename(output_dir: Path, rel_path: Path, metadata: Metadata) -> Path:
     output_file = change_suffix(
         output_dir.joinpath(rel_path),
-        get_metadata_value(metadata, "source.default.suffix", metadata["suffix"]),
+        metadata.get("source.default.suffix", metadata["suffix"]),
         ".jsonl.zst",
     )
     return output_file
 
 
 def process_file(
-    metadata: dict[str, Any],
+    metadata: Metadata,
     propella_lookup_fn: Callable,
     source_file: Path,
     output_file: Path,
diff --git a/training_data_packer/utils/file.py b/training_data_packer/utils/file.py
index acc37fd..630579f 100644
--- a/training_data_packer/utils/file.py
+++ b/training_data_packer/utils/file.py
@@ -10,7 +10,7 @@
 from loguru import logger
 
 
-def find_files(source_dir: Path, suffix: dict, part: str = None) -> list[Path]:
+def find_files(source_dir: Path, suffix: str, part: None | str = None) -> list[Path]:
     """
     Returns all files, not hidden, under source_dir, following symlinks
     :param source_dir: Dir to find files in
diff --git a/training_data_packer/utils/metadata.py b/training_data_packer/utils/metadata.py
index 05c929c..4818db9 100644
--- a/training_data_packer/utils/metadata.py
+++ b/training_data_packer/utils/metadata.py
@@ -1,4 +1,5 @@
 import hashlib
+from collections import UserDict
 from pathlib import Path
 from typing import Any
 
@@ -9,7 +10,25 @@
 from training_data_packer.utils.file import change_suffix
 
 
-def get_metadata_value(metadata: dict[str, Any], key: str, default: Any = None) -> Any:
+class Metadata(UserDict):
+    def __getitem__(self, key) -> Any:
+        key = key.replace("[", ".").replace("]", "")
+        return glom.glom(self.data, key)
+
+    def __setitem__(self, key, value) -> None:
+        key = key.replace("[", ".").replace("]", "")
+        return glom.assign(self.data, key, value)
+
+    def __delitem__(self, key):
+        key = key.replace("[", ".").replace("]", "")
+        return glom.delete(self.data, key)
+
+    def get(self, key, default=None):
+        key = key.replace("[", ".").replace("]", "")
+        return glom.glom(self.data, key, default=default)
+
+
+def get_metadata_value(metadata: Metadata, key: str, default: Any = None) -> Any:
     """
     Retrieve a specific value from the provided metadata structure using a
     dot-notation key for deep access.
@@ -22,10 +41,10 @@ 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.
     """
-    return glom.glom(metadata, key, default=default)
+    return metadata.get(key, default)
 
 
-def get_all_part_names(metadata: dict[str, Any]) -> list[str]:
+def get_all_part_names(metadata: Metadata) -> list[str]:
     """
     Returns all part names from metadata.
     :param metadata: Metadata dictionary.
@@ -36,12 +55,12 @@ def get_all_part_names(metadata: dict[str, Any]) -> list[str]:
     def _get_section_parts(section):
         section_keys = metadata[section].keys()
         section_parts = set(filter(lambda x: x not in reserved_part_names, section_keys))
-        input_src = get_metadata_value(metadata, f"{section}.default.input")
+        input_src = metadata.get(f"{section}.default.input")
         if input_src is not None:
             return section_parts.union(_get_section_parts(input_src))
         return section_parts
 
-    return sorted(_get_section_parts(get_metadata_value(metadata, "_internal.mode", "source")))
+    return sorted(_get_section_parts(metadata.get("_internal.mode", "source")))
 
 
 def get_shard_size_documents(part_config: dict[str, Any]) -> int:
@@ -62,7 +81,7 @@ def get_shard_size_documents(part_config: dict[str, Any]) -> int:
 
 
 def _get_pre_section_part(
-    metadata: dict[str, Any], src_file_name: Path, default_part_config, section_name: str
+    metadata: Metadata, src_file_name: Path, default_part_config, section_name: str
 ) -> tuple[None, None] | tuple[dict, str]:
     pre_section = metadata[section_name]
     for part in pre_section:
@@ -81,7 +100,7 @@ def _get_pre_section_part(
 
 
 def get_matching_part(
-    metadata: dict[str, Any], src_file_name: Path, section_name: str = "release"
+    metadata: Metadata, src_file_name: Path, section_name: str = "release"
 ) -> tuple[None, None] | tuple[dict, str]:
     """
     Returns matching part config and part name from metadata for given source file name.
@@ -110,7 +129,7 @@ def get_matching_part(
     return None, None
 
 
-def read_metadata(file_path: Path, log_content: bool = True) -> dict[str, Any]:
+def read_metadata(file_path: Path, log_content: bool = True) -> Metadata:
     """
     Reads metadata from file and returns it as dictionary.
     All field values are strings.
@@ -127,10 +146,10 @@ def read_metadata(file_path: Path, log_content: bool = True) -> dict[str, Any]:
         # BaseLoader to guarantee that the YAML parser will return unicode strings
         metadata = yaml.load(data, Loader=yaml.BaseLoader)
         metadata["_internal"] = {"collection_dir": file_path.parent, "sha256": sha256_data}
-        return metadata
+        return Metadata(metadata)
 
 
-def get_source_dir(metadata):
+def get_source_dir(metadata: Metadata) -> Path:
     mode = metadata["_internal"]["mode"]
     input_src = metadata[mode]["default"]["input"]
     if mode == "release" and "parallel" in metadata[mode]["default"]:
@@ -140,12 +159,12 @@ def get_source_dir(metadata):
     return metadata["_internal"]["collection_dir"].joinpath(input_src)
 
 
-def get_in_suffix(metadata: dict[str, Any], mode: str) -> str:
-    input_dir = get_metadata_value(metadata, f"{mode}.default.input", None)
-    return get_metadata_value(metadata, f"{input_dir}.default.suffix", metadata["suffix"])
+def get_in_suffix(metadata: Metadata, mode: str) -> str:
+    input_dir = metadata.get(f"{mode}.default.input", None)
+    return metadata.get(f"{input_dir}.default.suffix", metadata["suffix"])
 
 
-def calculate_file_path(src_file: Path, metadata: dict[str, Any], mode: str, process_dir: Path) -> Path:
+def calculate_file_path(src_file: Path, metadata: Metadata, mode: str, process_dir: Path) -> Path:
     input_suffix = get_in_suffix(metadata, mode)
     rel_file_path = src_file.relative_to(get_source_dir(metadata))
     out_suffix = ".jsonl.zst"
diff --git a/training_data_packer/utils/misc.py b/training_data_packer/utils/misc.py
index 7bb9010..f311c64 100644
--- a/training_data_packer/utils/misc.py
+++ b/training_data_packer/utils/misc.py
@@ -5,7 +5,7 @@
 from loguru import logger
 
 
-def hash_factory(hash_algo: str) -> Callable[str, str]:
+def hash_factory(hash_algo: str) -> Callable[[str], str]:
     match hash_algo.lower():
         case "sha256":
             return lambda x: hashlib.sha256(x.encode("utf-8")).hexdigest()