Merge pull request amueller#254 from icyblade/issue_253

add unicode support for stopwords
rgdk · Apr 19, 2017 · 35cdf67 · 35cdf67
2 parents 6a36626 + cfc86f6
commit 35cdf67
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/test/test_wordcloud.py b/test/test_wordcloud.py
@@ -246,3 +246,17 @@ def test_relative_scaling_zero():
     # non-regression test for non-integer font size
     wc = WordCloud(relative_scaling=0)
     wc.generate(THIS)
+
+
+def test_unicode_stopwords():
+    wc_unicode = WordCloud(stopwords=[u'Beautiful'])
+    try:
+        words_unicode = wc_unicode.process_text(unicode(THIS))
+    except NameError:  # PY3
+        words_unicode = wc_unicode.process_text(THIS)
+
+    wc_str = WordCloud(stopwords=['Beautiful'])
+    words_str = wc_str.process_text(str(THIS))
+
+    assert_true(words_unicode == words_str)
+
diff --git a/wordcloud/wordcloud.py b/wordcloud/wordcloud.py
@@ -503,7 +503,7 @@ def process_text(self, text):
         include all those things.
         """
 
-        stopwords = set(map(str.lower, self.stopwords))
+        stopwords = set([i.lower() for i in self.stopwords])
 
         flags = (re.UNICODE if sys.version < '3' and type(text) is unicode
                  else 0)