Generate iTOL "color strip" annotation file from pandas series

itol_colorstrip.py

import numpy as np
import pandas as pd
from ete3 import PhyloTree
from itertools import cycle

def generate_colorstrip(labels: pd.Series, save_file: str, palette: list = None) -> None:
    # labels must be a pd.Series whose index are leaf names and values are their labels
    if palette is None:
        palette = ["#4E79A7", "#A0CBE8", "#F28E2B", "#FFBE7D", "#59A14F",
              "#8CD17D", "#B6992D", "#F1CE63", "#499894", "#86BDB6",
              "#E15759", "#FF9D9A", "#79706E", "#BAB0AC", "#B37295",
              "#FABFD2", "#B07AA1", "#D4A6C8", "#9D7660", "#D7B5A6"]
    categories = labels.value_counts().index.tolist()
    categories_colors = {category: color for category, color in zip(categories, cycle(palette))}
    df = labels.rename("label").rename_axis("ids").reset_index()
    df["color"] = df["label"].map(categories_colors)
    output = """DATASET_COLORSTRIP\n
SEPARATOR COMMA\n
DATASET_LABEL,label_colorstrip\n
DATA\n"""+df.loc[:,["ids", "color", "label"]].to_csv(header=False, index=False)
    
    with open(save_file, "w") as f:
        f.write(output)

labels = pd.Series(["type A", "type B", "type B", "type C"], index=["sample 1", "sample 2", "sample 3", "sample 4"])
generate_colorstrip(labels, "colorstrip.txt")