Treemap with variable # of categories (Plotly)

I am working on a dashboard that uses a dataset for a treemap that looks like the following:

import pandas as pd
import plotly.express as px

col1 = ['United States','Mexico','Canada','Argentina','France','Portugal','Italy','Spain','Brazil','Peru']
col2 = [4,6,6,9,10,3,20,15,16,2]

df = pd.DataFrame({
    'country': col1,
    'count': col2
})

I will be using a filter that will change the countries and counts. Here is some code that will produce the treemap:

tree_fig = px.treemap(
    df, 
    path = ['country'],
    values = 'count',
    template='plotly_dark',
    color = 'country',
    color_discrete_map={
        df['country'][0]: '#626ffb', 
        df['country'][1]: '#b064fc', 
        df['country'][2]: '#ef563b',
        df['country'][3]: '#f45498',
        df['country'][4]: '#ff94fc',
        df['country'][5]: '#a8f064',
        df['country'][6]: '#24cce6',
        df['country'][7]: '#ffa45c',
        df['country'][8]: '#00cc96',
        df['country'][9]: '#fff000'

    }   
)
tree_fig.update_traces(
    hovertemplate='Count=%{value}'
)
tree_fig.show()

This all works, but when using the filters there will be situations where I have fewer than 10 countries in the dataframe. How can I rewrite the color_discrete_map argument to account for a varying number of categories.

You can do something like this:

# create a function which creates colour maps
def colour_map_creator(df, col):
    colour_map = {}
    for i in range(len(df[col])):
        colour_map[df[col][i]] = px.colors.qualitative.Plotly[i]
    return colour_map

Then replace the dictionary which is the value of color_discrete_map with color_map_creator(df, 'country').