Renumber taxids in taxdumps

fix_taxdump.py

#!/usr/bin/env python

import argparse
from pathlib import Path


def main():
    parser = argparse.ArgumentParser(
        description="Renumber taxids in taxdumps."
    )
    parser.add_argument(
        "TAXDUMP_DIRECTORY",
        help="""
            Path to the taxdump directory containing names.dmp, nodes.dmp, and
            optionally taxid.map, merged.dmp, and delnodes.dmp files.
        """,
    )

    args = parser.parse_args()
    taxdump_dir = Path(args.TAXDUMP_DIRECTORY)

    taxid_mapping_dict = {}
    count = 1

    with open(taxdump_dir.joinpath("names.dmp")) as fin:
        for line in fin:
            taxid = line.split("\t")[0]
            taxid_mapping_dict[taxid] = str(count)
            count += 1

    with (
        open(taxdump_dir.joinpath("names.dmp")) as fin,
        open(taxdump_dir.joinpath("names.dmp.new"), "w") as fout,
    ):
        for line in fin:
            line = line.strip().split("\t")
            line[0] = taxid_mapping_dict[line[0]]
            line = "\t".join(line)
            fout.write(f"{line}\n")

    with (
        open(taxdump_dir.joinpath("nodes.dmp")) as fin,
        open(taxdump_dir.joinpath("nodes.dmp.new"), "w") as fout,
    ):
        for line in fin:
            line = line.strip().split("\t")
            line[0] = taxid_mapping_dict[line[0]]
            line[2] = taxid_mapping_dict[line[2]]
            line = "\t".join(line)
            fout.write(f"{line}\n")

    if taxdump_dir.joinpath("taxid.map").exists():
        with (
            open(taxdump_dir.joinpath("taxid.map")) as fin,
            open(taxdump_dir.joinpath("taxid.map.new"), "w") as fout,
        ):
            for line in fin:
                line = line.strip().split("\t")
                line[1] = taxid_mapping_dict[line[1]]
                line = "\t".join(line)
                fout.write(f"{line}\n")

    if taxdump_dir.joinpath("merged.dmp").exists():
        with (
            open(taxdump_dir.joinpath("merged.dmp")) as fin,
            open(taxdump_dir.joinpath("merged.dmp.new"), "w") as fout,
        ):
            for line in fin:
                line = line.strip().split("\t")
                line[0] = taxid_mapping_dict[line[0]]
                line[2] = taxid_mapping_dict[line[2]]
                line = "\t".join(line)
                fout.write(f"{line}\n")

    if taxdump_dir.joinpath("delnodes.dmp").exists():
        with (
            open(taxdump_dir.joinpath("delnodes.dmp")) as fin,
            open(taxdump_dir.joinpath("delnodes.dmp.new"), "w") as fout,
        ):
            for line in fin:
                line = line.strip().split("\t")
                line[0] = taxid_mapping_dict[line[0]]
                line = "\t".join(line)
                fout.write(f"{line}\n")

    taxdump_dir.joinpath("names.dmp").unlink()
    taxdump_dir.joinpath("nodes.dmp").unlink()

    taxdump_dir.joinpath("nodes.dmp.new").rename(taxdump_dir.joinpath("nodes.dmp"))
    taxdump_dir.joinpath("names.dmp.new").rename(taxdump_dir.joinpath("names.dmp"))

    if taxdump_dir.joinpath("taxid.map").exists():
        taxdump_dir.joinpath("taxid.map").unlink()
        taxdump_dir.joinpath("taxid.map.new").rename(taxdump_dir.joinpath("taxid.map"))

    if taxdump_dir.joinpath("merged.dmp").exists():
        taxdump_dir.joinpath("merged.dmp").unlink()
        taxdump_dir.joinpath("merged.dmp.new").rename(
            taxdump_dir.joinpath("merged.dmp")
        )

    if taxdump_dir.joinpath("delnodes.dmp").exists():
        taxdump_dir.joinpath("delnodes.dmp").unlink()
        taxdump_dir.joinpath("delnodes.dmp.new").rename(
            taxdump_dir.joinpath("delnodes.dmp")
        )


if __name__ == "__main__":
    main()