Allow setting font path and font color by frequency

Author: naslundx

.pylintrc

@@ -418,7 +418,7 @@ max-spelling-suggestions=2
 [DESIGN]
 
 # Maximum number of arguments for function / method
-max-args=10
+max-args=15
 
 # Maximum number of locals for function / method body
 max-locals=25

examples/sherlock.py

@@ -1,3 +1,4 @@
+import math
 from pathlib import Path
 
 from wrdcld import make_word_cloud
@@ -7,9 +8,17 @@
 all_words = [word.strip(" \n,.!?:-&\"'[]") for word in contents.split(" ")]
 all_words = [word for word in all_words if word]
 
-make_word_cloud(
+
+def color_func(frequency):
+    x = int((1 - frequency) * 200)
+    return (x, x, x)
+
+
+img = make_word_cloud(
     all_words=all_words,
-    font_color=(0, 0, 0),
+    font_color_func=color_func,
     background_color=(255, 255, 255),
     minimum_font_size=5,
-).show()
+)
+
+img.show()

wrdcld/__init__.py

@@ -1,20 +1,24 @@
 import math
 from collections import Counter
 from collections.abc import Callable
+from pathlib import Path
 
 from .font import FontWrapper
 from .image import ImageWrapper
 from .main import fill_next_word
 from .rectangle import Rectangle
+from .util import Color, get_repo_root
 
 
 # pylint: disable=unused-argument
 def make_word_cloud(
     all_words: list[str],
     width: int = 500,
     height: int = 500,
-    font_color: tuple[int, int, int] = (255, 255, 0),
-    background_color: tuple[int, int, int] = (73, 109, 137),
+    font_path: Path = None,
+    font_color: Color = (255, 255, 0),
+    font_color_func: Callable[[float], Color] = None,
+    background_color: Color = (73, 109, 137),
     minimum_font_size: int = 1,
     maximum_font_size: int = 100,
     word_padding: int = 0,  # TODO
@@ -28,24 +32,32 @@ def make_word_cloud(
     assert (
         0 < minimum_font_size < maximum_font_size
     ), "Invalid font sizes, must be positive (in pixels)"
+    assert (
+        font_color is not None or font_color_func is not None
+    ), "Must specify a fixed font color or function"
 
     # Create a new image and font
+    font_path = font_path or (get_repo_root() / "fonts" / "OpenSans-Regular.ttf")
+    font_color_func = font_color_func or (lambda _: font_color)
     image = ImageWrapper(width, height, background_color)
-    font = FontWrapper(color=font_color, size=maximum_font_size)
+    font = FontWrapper(
+        path=font_path, color_func=font_color_func, size=maximum_font_size
+    )
 
     # Handle data
     word_counts = Counter(all_words)
     _, first_count = word_counts.most_common(1)[0]
 
     available_rectangles = [Rectangle(width=width, height=height, x=0, y=0)]
     for word, count in word_counts.most_common():
-        required_font_size = maximum_font_size * scaling_func(count / first_count)
+        frequency = count / first_count
+        required_font_size = maximum_font_size * scaling_func(frequency)
 
         if required_font_size < minimum_font_size:
             break
 
         available_rectangles = fill_next_word(
-            word, available_rectangles, image, font[required_font_size]
+            word, available_rectangles, image, font[required_font_size], frequency
         )
 
     return image.img

wrdcld/font.py

@@ -1,3 +1,4 @@
+from collections.abc import Callable
 from dataclasses import dataclass, replace
 from functools import lru_cache
 from pathlib import Path
@@ -6,14 +7,17 @@
 
 from .image import ImageWrapper
 from .rectangle import Rectangle
-from .util import get_repo_root
+from .util import Color
 
 
 @dataclass(frozen=True)
 class FontWrapper:
-    path: Path = get_repo_root() / "fonts" / "OpenSans-Regular.ttf"
+    color_func: Callable[[float], Color]
+    path: Path
     size: int = 1
-    color: tuple[int, int, int] = (255, 255, 0)
+
+    def color(self, frequency: float) -> Color:
+        return self.color_func(frequency)
 
     @lru_cache(maxsize=1024)
     def get(self):
@@ -39,6 +43,7 @@ def draw_text(
     rectangle: Rectangle,
     word: str,
     font: FontWrapper,
+    frequency: float,
     rotate=False,
 ):
     """
@@ -52,7 +57,10 @@ def draw_text(
         text_draw = ImageDraw.Draw(text_image)
 
         text_draw.text(
-            (-text_bbox.x, -text_bbox.y), word, font=font.get(), fill=font.color
+            (-text_bbox.x, -text_bbox.y),
+            word,
+            font=font.get(),
+            fill=font.color(frequency),
         )
         rotated_text_image = text_image.rotate(90, expand=True)
         image.img.paste(rotated_text_image, rectangle.xy)
@@ -62,5 +70,5 @@ def draw_text(
             (rectangle.x - text_bbox.x, rectangle.y - text_bbox.y),
             word,
             font=font.get(),
-            fill=font.color,
+            fill=font.color(frequency),
         )

wrdcld/image.py

@@ -1,8 +1,10 @@
 from PIL import Image, ImageDraw
 
+from .util import Color
+
 
 class ImageWrapper:
-    def __init__(self, width: int, height: int, background_color: tuple[int, int, int]):
+    def __init__(self, width: int, height: int, background_color: Color):
         self.width = width
         self.height = height
         self.background_color = background_color

wrdcld/main.py

@@ -16,6 +16,7 @@ def _fill(
     word_length: int,
     word: str,
     font: FontWrapper,
+    frequency: float,
     rotate: bool = False,
 ):
     word_height = font.size
@@ -47,12 +48,12 @@ def _fill(
             height=word_length,
         )
 
-    draw_text(image, text_rectangle, word, font, rotate=rotate)
+    draw_text(image, text_rectangle, word, font, frequency, rotate=rotate)
 
     return text_rectangle
 
 
-def fill_next_word(word, available_rectangles, image, font):
+def fill_next_word(word, available_rectangles, image, font, frequency):
     word_length = font.get_length_of_word(word)
 
     suitable_horizontal_rectangles = [
@@ -96,7 +97,9 @@ def fill_next_word(word, available_rectangles, image, font):
     if option == "horizontal":
         available_rectangles.remove(horizontal_option)
         chosen_rectangle = horizontal_option
-        text_rectangle = _fill(chosen_rectangle, image, word_length, word, font)
+        text_rectangle = _fill(
+            chosen_rectangle, image, word_length, word, font, frequency
+        )
 
     else:
         available_rectangles.remove(vertical_option)
@@ -107,6 +110,7 @@ def fill_next_word(word, available_rectangles, image, font):
             word_length,
             word,
             font,
+            frequency,
             rotate=True,
         )
 

wrdcld/util.py

@@ -1,5 +1,9 @@
+from typing import TypeAlias
 from pathlib import Path
 
 
+Color: TypeAlias = tuple[int, int, int]
+
+
 def get_repo_root():
     return Path(__file__).parent.parent