Commit 7e3f1162 authored by vkolmakov's avatar vkolmakov

Refactor into class and provide another class that uses exact colors

parent f609bfeb
......@@ -10,38 +10,66 @@ from wordcloud import (WordCloud, get_single_color_func)
import matplotlib.pyplot as plt
def get_grouped_color_func(colors_to_words, default_color):
"""Create a color function which assigns colors to certain words
based on the color to words mapping
class SimpleGroupedColorFunc(object):
"""Create a color function object which assigns EXACT colors
to certain words based on the color to words mapping
Parameters
----------
colors_to_words : dict(str -> list(str))
color_to_words : dict(str -> list(str))
A dictionary that maps a color to the list of words.
default_color : str
Color that will be assigned to a word that's not a member
of any value from colors_to_words.
of any value from color_to_words.
"""
color_func_to_words = [(get_single_color_func(color), set(words))
for (color, words) in colors_to_words.items()]
default_color_func = get_single_color_func(default_color)
def __init__(self, color_to_words, default_color):
self.word_to_color = {word: color
for (color, words) in color_to_words.items()
for word in words}
self.default_color = default_color
def grouped_color_func(word=None, font_size=None, position=None,
orientation=None, font_path=None, random_state=None, **kwargs):
def __call__(self, word, **kwargs):
return self.word_to_color.get(word, self.default_color)
class GroupedColorFunc(object):
"""Create a color function object which assigns DIFFERENT SHADES of
specified colors to certain words based on the color to words mapping.
Uses wordcloud.get_single_color_func
Parameters
----------
color_to_words : dict(str -> list(str))
A dictionary that maps a color to the list of words.
default_color : str
Color that will be assigned to a word that's not a member
of any value from color_to_words.
"""
def __init__(self, color_to_words, default_color):
self.color_func_to_words = [(get_single_color_func(color), set(words))
for (color, words) in color_to_words.items()]
self.default_color_func = get_single_color_func(default_color)
def get_color_func(self, word):
"""Returns a single_color_func associated with the word"""
try:
color_func = next(color_func
for (color_func, words) in color_func_to_words
for (color_func, words) in self.color_func_to_words
if word in words)
except StopIteration:
color_func = default_color_func
color_func = self.default_color_func
return color_func(word, font_size, position,
orientation, font_path, random_state, **kwargs)
return color_func
return grouped_color_func
def __call__(self, word, **kwargs):
return self.get_color_func(word)(word, **kwargs)
text = """The Zen of Python, by Tim Peters
......@@ -68,7 +96,7 @@ Namespaces are one honking great idea -- let's do more of those!"""
# Since the text is small collocations are turned off and text is lower-cased
wc = WordCloud(collocations=False).generate(text.lower())
colors_to_words = {
color_to_words = {
# words below will be colored with a green single color function
'#00ff00': ['beautiful', 'explicit', 'simple', 'sparse',
'readability', 'rules', 'practicality',
......@@ -79,12 +107,15 @@ colors_to_words = {
'guess', 'hard']
}
# Words that are not in any of the colors_to_words values
# Words that are not in any of the color_to_words values
# will be colored with a grey single color function
default_color = 'grey'
# Create a color function
grouped_color_func = get_grouped_color_func(colors_to_words, default_color)
# Create a color function with single tone
# grouped_color_func = SimpleGroupedColorFunc(color_to_words, default_color)
# Create a color function with multiple tones
grouped_color_func = GroupedColorFunc(color_to_words, default_color)
# Apply our color function
wc.recolor(color_func=grouped_color_func)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment