diff --git a/tests/test_compression.py b/tests/test_compression.py index 2fb50c5..709f2cf 100644 --- a/tests/test_compression.py +++ b/tests/test_compression.py @@ -113,6 +113,19 @@ def test_small_msprime_migration(self): self.assertGreater(ts.num_trees, 2) self.verify(ts) + def test_small_msprime_top_level_metadata(self): + ts = msprime.simulate(10, recombination_rate=2, mutation_rate=2, random_seed=2) + self.assertGreater(ts.num_sites, 2) + self.assertGreater(ts.num_trees, 2) + tables = ts.dump_tables() + top_level_schema = { + 'codec': 'json', + 'properties': {'my_int': {'type': 'integer'}} + } + tables.metadata_schema = tskit.MetadataSchema(top_level_schema) + tables.metadata = {"my_int": 1234} + self.verify(tables.tree_sequence()) + def test_small_msprime_individuals_metadata(self): ts = msprime.simulate(10, recombination_rate=1, mutation_rate=2, random_seed=2) self.assertGreater(ts.num_sites, 2) diff --git a/tszip/compression.py b/tszip/compression.py index 759264b..41d430a 100644 --- a/tszip/compression.py +++ b/tszip/compression.py @@ -172,6 +172,9 @@ def compress_zarr(ts, root, variants_only=False): root.attrs["format_version"] = FORMAT_VERSION root.attrs["sequence_length"] = tables.sequence_length root.attrs["provenance"] = provenance_dict + if tables.metadata_schema.schema is not None: + root.attrs["metadata_schema"] = tables.metadata_schema.schema + root.attrs["metadata"] = tables.metadata columns = [ Column("coordinates", coordinates), @@ -280,6 +283,9 @@ def load_zarr(path): def decompress_zarr(root): tables = tskit.TableCollection(root.attrs["sequence_length"]) coordinates = root["coordinates"][:] + if "metadata_schema" in root.attrs and "metadata" in root.attrs: + tables.metadata_schema = tskit.MetadataSchema(root.attrs["metadata_schema"]) + tables.metadata = root.attrs["metadata"] tables.individuals.set_columns( flags=root["individuals/flags"],