Script to add node data as branch label

add-branch-labels.py

import argparse
import json


def extract_attribute(node_data, attribute: str) -> dict[str, str]:
    data = {}
    for name, node in node_data["nodes"].items():
        if attribute in node:
            data[name] = node[attribute]
    return data


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Add extra branch labels",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument(
        "--input-tree",
        type=str,
        metavar="JSON",
        required=True,
        help="input Auspice JSON",
    )
    parser.add_argument(
        "--input-node-data", type=str, required=True, help="node data file"
    )
    parser.add_argument(
        "--attribute",
        type=str,
        required=True,
        help="Attribute from node data to add as branch label",
    )
    parser.add_argument(
        "--output",
        type=str,
        metavar="JSON",
        required=True,
        help="output Auspice JSON",
    )
    args = parser.parse_args()

    with open(args.input_tree, "r") as f:
        auspice_json = json.load(f)

    with open(args.input_node_data, "r") as f:
        node_data = extract_attribute(json.load(f), args.attribute)

    def attach_labels(n: dict[str,dict]):  # closure
        if n["name"] in node_data:
            if "branch_attrs" not in n:
                n["branch_attrs"] = {}
            if "labels" not in n["branch_attrs"]:
                n["branch_attrs"]["labels"] = {}
            n["branch_attrs"]["labels"][
                args.attribute
            ] = node_data[n["name"]]

        if "children" in n:
            for c in n["children"]:
                attach_labels(c)

    attach_labels(auspice_json["tree"])

    with open(args.output, "w") as f:
        json.dump(auspice_json, f, indent=2)