WordCloud Custom Colors

from wordcloud import WordCloud

text = 'i solemnly swear i am up to no good mischief managed'

#  определяем цвета – списки positive и negative внутри
# если положительно – красный, если отрицательно – черный, иначе – серый

def my_color_func(word, font_size, position, orientation, random_state=None,
                    **kwargs):
    positive = ['managed', 'good', 'up']
    negative = ['no', 'mischief']
    if word in positive:
        col = 'red'
    elif word in negative:
        col = 'black'
    else:
        col = 'grey'
    return col

# строим облако с цветами по умолчанию

wordcloud = WordCloud(background_color = "white", 
                     max_words = 50).generate(text)
                     
# сохраняем набор цветов по умолчанию в массив default_colors
# используем recolor и туда подставляем свою функцию
# должно все покраситься так, как задали в функции

default_colors = wordcloud.to_array()
plt.imshow(wordcloud.recolor(color_func=my_color_func, random_state=3),
           interpolation="bilinear")
plt.axis("off")
plt.show()