test: add tests of Python consuming JSON metadata from rust (#749)

Author: molpopgen

.github/workflows/python.yml

@@ -0,0 +1,47 @@
+on:
+  push:
+    branches: [main, dev]
+  pull_request:
+
+name: Python data round trips
+
+jobs:
+  test-metadata:
+    name: Test Python metadata round trips
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-24.04]
+        rust:
+          - stable
+        python: [ "3.13" ]
+    steps:
+      - name: Cancel Previous Runs
+        uses: styfle/[email protected]
+        with:
+          access_token: ${{ secrets.GITHUB_TOKEN }}
+
+      - uses: actions/[email protected]
+        with:
+          token: ${{ secrets.GITHUB_TOKEN }}
+          submodules: recursive
+      - uses: dtolnay/rust-toolchain@v1
+        with:
+          toolchain: ${{ matrix.rust }}
+      - uses: Swatinem/[email protected]
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          activate-environment: true
+          version: "latest"
+          python-version: ${{ matrix.python }}
+      - name: run JSON metadata example
+        run: |
+          cargo run --example json_metadata --features derive
+      - name: setup Python and run tests
+        run: |
+          uv venv -p ${{ matrix.python }}
+          source .venv/bin/activate
+          uv pip install -r python/requirements_locked_3_13.txt
+          python -m pytest python
+

Cargo.toml

@@ -56,3 +56,7 @@ rustdoc-args = ["--cfg", "doc_cfg"]
 # Not run during tests
 [[example]]
 name = "tree_traversals"
+
+[[example]]
+name = "json_metadata"
+required-features = ["derive"]

examples/json_metadata.rs

@@ -0,0 +1,51 @@
+#[derive(serde::Serialize, serde::Deserialize, tskit::metadata::MutationMetadata)]
+#[serializer("serde_json")]
+struct MutationMetadata {
+    effect_size: f64,
+    dominance: f64,
+}
+
+#[derive(serde::Serialize, serde::Deserialize, tskit::metadata::IndividualMetadata)]
+#[serializer("serde_json")]
+struct IndividualMetadata {
+    name: String,
+    phenotypes: Vec<i32>,
+}
+
+fn main() {
+    let ts = make_treeseq().unwrap();
+    ts.dump("with_json_metadata.trees", 0).unwrap();
+}
+
+fn make_tables() -> anyhow::Result<tskit::TableCollection> {
+    let mut tables = tskit::TableCollection::new(100.0)?;
+    let pop0 = tables.add_population()?;
+    let ind0 = tables.add_individual_with_metadata(
+        0,
+        None,
+        None,
+        &IndividualMetadata {
+            name: "Jerome".to_string(),
+            phenotypes: vec![0, 1, 2, 0],
+        },
+    )?;
+    let node0 = tables.add_node(tskit::NodeFlags::new_sample(), 0.0, pop0, ind0)?;
+    let site0 = tables.add_site(50.0, Some("A".as_bytes()))?;
+    let _ = tables.add_mutation_with_metadata(
+        site0,
+        node0,
+        tskit::MutationId::NULL,
+        1.0,
+        Some("G".as_bytes()),
+        &MutationMetadata {
+            effect_size: -1e-3,
+            dominance: 0.1,
+        },
+    )?;
+    tables.build_index()?;
+    Ok(tables)
+}
+
+fn make_treeseq() -> anyhow::Result<tskit::TreeSequence> {
+    Ok(make_tables()?.tree_sequence(0)?)
+}

python/requirements.txt

@@ -0,0 +1,2 @@
+tskit>=0.6.3
+pytest

python/requirements_locked_3_13.txt

@@ -0,0 +1,30 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile -p 3.13 requirements.txt
+attrs==25.3.0
+    # via
+    #   jsonschema
+    #   referencing
+iniconfig==2.1.0
+    # via pytest
+jsonschema==4.23.0
+    # via tskit
+jsonschema-specifications==2025.4.1
+    # via jsonschema
+numpy==2.2.5
+    # via tskit
+packaging==25.0
+    # via pytest
+pluggy==1.6.0
+    # via pytest
+pytest==8.3.5
+    # via -r requirements.txt
+referencing==0.36.2
+    # via
+    #   jsonschema
+    #   jsonschema-specifications
+rpds-py==0.25.0
+    # via
+    #   jsonschema
+    #   referencing
+tskit==0.6.3
+    # via -r requirements.txt

python/test_json_metadata.py

@@ -0,0 +1,71 @@
+import tskit
+import numpy as np
+
+
+def setup_ts_without_schema():
+    ts = tskit.TreeSequence.load("with_json_metadata.trees")
+    return ts
+
+
+def setup_ts_with_schema():
+    ts = setup_ts_without_schema()
+    tables = ts.tables
+    tables.individuals.metadata_schema = tskit.metadata.MetadataSchema(
+        {
+            "codec": "json",
+            "type": "object",
+            "name": "Individual metadata",
+            "properties": {"name": {"type": "string"},
+                           "phenotypes": {"type": "array"}},
+            "additionalProperties": False,
+        })
+    tables.mutations.metadata_schema = tskit.metadata.MetadataSchema(
+        {
+            "codec": "json",
+            "type": "object",
+            "name": "Individual metadata",
+            "properties": {"effect_size": {"type": "number"},
+                           "dominance": {"type": "number"}},
+            "additionalProperties": False,
+        })
+    return tables.tree_sequence()
+
+
+def test_individual_metadata():
+    # NOTE: the assertions here rely on knowing
+    # what examples/json_metadata.rs put into the
+    # metadata!
+    ts = setup_ts_with_schema()
+    md = ts.individual(0).metadata
+    assert md["name"] == "Jerome"
+    assert md["phenotypes"] == [0, 1, 2, 0]
+
+
+def test_individual_metadata_without_schema():
+    # NOTE: the assertions here rely on knowing
+    # what examples/json_metadata.rs put into the
+    # metadata!
+    ts = setup_ts_without_schema()
+    md = eval(ts.individual(0).metadata)
+    assert md["name"] == "Jerome"
+    assert md["phenotypes"] == [0, 1, 2, 0]
+
+
+def test_mutation_metadata():
+    # NOTE: the assertions here rely on knowing
+    # what examples/json_metadata.rs put into the
+    # metadata!
+    ts = setup_ts_with_schema()
+    md = ts.mutation(0).metadata
+    assert np.isclose(md["effect_size"], -1e-3)
+    assert np.isclose(md["dominance"], 0.1)
+
+
+def test_mutation_metadata_without_schema():
+    # NOTE: the assertions here rely on knowing
+    # what examples/json_metadata.rs put into the
+    # metadata!
+    ts = setup_ts_without_schema()
+    md = eval(ts.mutation(0).metadata)
+    assert np.isclose(md["effect_size"], -1e-3)
+    assert np.isclose(md["dominance"], 0.1)