Allow Counter as all_words type

Author: naslundx

tests/test_wrdcld.py

@@ -1,3 +1,4 @@
+from collections import Counter
 from string import ascii_letters, digits, punctuation
 from unittest import TestCase
 
@@ -56,8 +57,9 @@ def test_make_word_cloud(self, words: list[str]):
     @settings(deadline=None)
     @given(words=words_with_repeats_strategy(), seed=st.integers())
     def test_word_cloud_reproducibility(self, words: list[str], seed: int):
+        counter = Counter(words)
 
         word_cloud_1 = make_word_cloud(words, seed=seed)
-        word_cloud_2 = make_word_cloud(words, seed=seed)
+        word_cloud_2 = make_word_cloud(counter, seed=seed)
 
         self.assertTrue(_two_images_are_equal(word_cloud_1, word_cloud_2))

wrdcld/__init__.py

@@ -16,7 +16,7 @@
 
 # pylint: disable=unused-argument
 def make_word_cloud(
-    all_words: list[str],
+    all_words: list[str] | Counter,
     width: int = 500,
     height: int = 500,
     font_path: Path | None = None,
@@ -32,6 +32,7 @@ def make_word_cloud(
 ):
     if seed is not None:
         random.seed(seed)
+
     # Asserts
     assert len(all_words) > 0, "No words in list"
     assert width > 0, "Width must be a positive number (in pixels)"
@@ -42,9 +43,13 @@ def make_word_cloud(
     assert (
         font_color is not None or font_color_func is not None
     ), "Must specify a fixed font color or function"
+    assert isinstance(all_words, (list, Counter)), "Word supply must be list or Counter"
 
     # Count the words
-    word_counts = Counter(all_words)
+    if isinstance(all_words, list):
+        word_counts = Counter(all_words)
+    else:
+        word_counts = all_words
     first_word, first_count = word_counts.most_common(1)[0]
 
     # Create a new image and font