load_in_ITOL.py

#!/usr/bin/env python3

import argparse
import os.path

from src.utils_argparse import MyFormatter
from src.utils_ITOL import create_style_file_with_simulation, upload_to_itol, reroot_tree_for_visualisation

if __name__ == "__main__":
    ## \cond
    parser = argparse.ArgumentParser(description="BLABLA", formatter_class=MyFormatter)
    parser.add_argument("project_name", help="Project name. All generated files will be put in DATA_generated/<PROJECT_NAME>")
    parser.add_argument("--ITOL_info", nargs=2, metavar=('batch_ID', 'project_name'),
                               help=" ITOL Batch ID and EXISTING project name to immediately load the phylum in ITOL")

    args = parser.parse_args()

    path_project = os.path.join(os.path.dirname(os.path.realpath(__file__)), "DATA_generated", args.project_name)

    if os.path.isdir(os.path.join(path_project, "iqtree", "ancestral")):
        tree_files = (file for file in os.listdir(os.path.join(path_project, "iqtree", "ancestral")) if file.endswith(".treefile"))
        file_loaded = False
        for file in tree_files:
            if args.ITOL_info:
                file_loaded = True
                rooted_tree_file = os.path.join(path_project, "iqtree", "ancestral", "ITOL_" + file.split('.')[0] + ".tree")
                reroot_tree_for_visualisation(os.path.join(path_project, "iqtree", "ancestral", file), rooted_tree_file)
                upload_to_itol(rooted_tree_file, dataset_paths=[], tree_name=args.project_name + "_" + file.split(".")[0], tree_description=None, project_name=args.ITOL_info[1],
                               upload_id=args.ITOL_info[0], success_message= "--> " + file.split(".")[0] + " - Real tree with internal node name succesfully loaded." )
                print()
        if file_loaded:
            print("----- Activate 'Node IDs' display in the 'Advanced' panel to see node names-----\n")

    if os.path.isdir(os.path.join(path_project, "iqtree", "with_simulation")):
        tree_files = (file for file in os.listdir(os.path.join(path_project, "iqtree", "with_simulation")) if file.endswith(".treefile"))
        for file in tree_files:
            new_tree, annotation = create_style_file_with_simulation(os.path.join(path_project, "iqtree", "with_simulation", file), os.path.join(path_project, "iqtree", "with_simulation", "ITOL_with_simulation_" + file.split(".")[0]))
            if args.ITOL_info:
                upload_to_itol(new_tree, dataset_paths=[annotation], tree_name=args.project_name + "_simulation_" + file.split(".")[0], tree_description=None, project_name=args.ITOL_info[1],
                               upload_id=args.ITOL_info[0], success_message= "--> " + file.split(".")[0] + " - Tree with simulation sequences successfully uploaded")
                print()
            else:
                print("--> " + file.split(".")[0] + " - Tree with simulation sequences successfully created :")
                print("    " + new_tree)
                print("    " + annotation)
    if args.ITOL_info is None:
        print("\nProvide ITOL info (batch ID and project name) with option --ITOL_info to automatically load trees on ITOL")