Merge pull request #111 from pendragon-/master

Add a command line interface

continuous_integration/install.sh

@@ -28,7 +28,7 @@ if [[ "$DISTRIB" == "conda" ]]; then
 
     # Configure the conda environment and put it in the path using the
     # provided versions
-    conda create -n testenv --yes python=$PYTHON_VERSION pip nose \
+    conda create -n testenv --yes python=$PYTHON_VERSION pip nose mock \
         numpy=$NUMPY_VERSION
     source activate testenv
 
@@ -37,7 +37,7 @@ elif [[ "$DISTRIB" == "ubuntu" ]]; then
     # Use standard ubuntu packages in their default version
     virtualenv --system-site-packages testvenv
     source testvenv/bin/activate
-    pip install nose
+    pip install nose mock
 fi
 
 pip install -r requirements.txt

setup.py

@@ -9,6 +9,7 @@ setup(
     license='MIT',
     ext_modules=[Extension("wordcloud.query_integral_image",
                            ["wordcloud/query_integral_image.c"])],
+    scripts=['wordcloud/wordcloud_cli.py'],
     packages=['wordcloud'],
     package_data={'wordcloud': ['stopwords', 'DroidSansMono.ttf']}
 )

test/test_wordcloud.py

-from wordcloud import WordCloud
+from wordcloud import WordCloud, get_single_color_func
 import numpy as np
+from random import Random
 from nose.tools import assert_equal, assert_greater, assert_true, assert_raises
 from numpy.testing import assert_array_equal
 from PIL import Image
@@ -131,6 +132,35 @@ def test_mask():
     assert_greater(wc_array[mask == 0].sum(), 10000)
 
 
+def test_single_color_func():
+    # test single color function for different color formats
+    random = Random(42)
+
+    red_function = get_single_color_func('red')
+    assert_equal(red_function(random_state=random), 'rgb(181, 0, 0)')
+
+    hex_function = get_single_color_func('#00b4d2')
+    assert_equal(hex_function(random_state=random), 'rgb(0, 48, 56)')
+
+    rgb_function = get_single_color_func('rgb(0,255,0)')
+    assert_equal(rgb_function(random_state=random), 'rgb(0, 107, 0)')
+
+    rgb_perc_fun = get_single_color_func('rgb(80%,60%,40%)')
+    assert_equal(rgb_perc_fun(random_state=random), 'rgb(97, 72, 48)')
+
+    hsl_function = get_single_color_func('hsl(0,100%,50%)')
+    assert_equal(hsl_function(random_state=random), 'rgb(201, 0, 0)')
+
+
+def test_single_color_func_grey():
+    # grey is special as it's a corner case
+    random = Random(42)
+
+    red_function = get_single_color_func('darkgrey')
+    assert_equal(red_function(random_state=random), 'rgb(181, 181, 181)')
+    assert_equal(red_function(random_state=random), 'rgb(56, 56, 56)')
+
+
 def check_parameters():
     # check that parameters are actually used
     pass

test/test_wordcloud_cli.py

+import argparse
+import inspect
+from wordcloud import wordcloud_cli as cli
+import wordcloud as wc
+from collections import namedtuple
+from mock import patch, MagicMock
+from nose.tools import assert_equal, assert_true, assert_in
+
+from tempfile import NamedTemporaryFile
+
+temp = NamedTemporaryFile()
+ArgOption = namedtuple('ArgOption', ['cli_name', 'init_name', 'pass_value', 'fail_value'])
+ARGUMENT_SPEC_TYPED = [
+    ArgOption(cli_name='width', init_name='width', pass_value=13, fail_value=1.),
+    ArgOption(cli_name='height', init_name='height', pass_value=15, fail_value=1.),
+    ArgOption(cli_name='margin', init_name='margin', pass_value=17, fail_value=1.),
+    ArgOption(cli_name='relative_scaling', init_name='relative_scaling', pass_value=1, fail_value='c')
+]
+ARGUMENT_SPEC_REMAINING = [
+    ArgOption(cli_name='stopwords', init_name='stopwords', pass_value=temp.name, fail_value=None),
+    ArgOption(cli_name='mask', init_name='mask', pass_value=temp.name, fail_value=None),
+    ArgOption(cli_name='fontfile', init_name='font_path', pass_value=temp.name, fail_value=None),
+    ArgOption(cli_name='color', init_name='color_func', pass_value='red', fail_value=None),
+    ArgOption(cli_name='background', init_name='background_color', pass_value='grey', fail_value=None)
+]
+
+def all_arguments():
+    arguments = []
+    arguments.extend(ARGUMENT_SPEC_TYPED)
+    arguments.extend(ARGUMENT_SPEC_REMAINING)
+    return arguments
+
+
+def test_argument_spec_matches_to_constructor_args():
+    args = argparse.Namespace()
+    for option in all_arguments():
+        setattr(args, option.init_name, option.pass_value)
+
+    supported_args = inspect.getargspec(wc.WordCloud.__init__).args
+    supported_args.remove('self')
+    for arg_name in vars(args).keys():
+        assert_in(arg_name, supported_args)
+
+
+def test_main_passes_arguments_through():
+    args = argparse.Namespace()
+    for option in all_arguments():
+        setattr(args, option.init_name, option.pass_value)
+    args.imagefile = NamedTemporaryFile()
+    args.text = 'some long text'
+
+    with patch('wordcloud.wordcloud_cli.wc.WordCloud', autospec=True) as mock_word_cloud:
+        instance = mock_word_cloud.return_value
+        instance.to_image.return_value = MagicMock()
+        cli.main(args)
+
+    posargs, kwargs = mock_word_cloud.call_args
+    for option in all_arguments():
+        assert_in(option.init_name, kwargs)
+
+
+def check_argument(name, result_name, value):
+    text = NamedTemporaryFile()
+
+    args = cli.parse_args(['--text', text.name, '--' + name, str(value)])
+    assert_in(result_name, vars(args))
+
+
+def check_argument_type(name, value):
+    text = NamedTemporaryFile()
+
+    try:
+        with patch('sys.stderr') as mock_stderr:
+            args = cli.parse_args(['--text', text.name, '--' + name, str(value)])
+        raise AssertionError('argument "{}" was accepted even though the type did not match'.format(name))
+    except SystemExit:
+        pass
+    except ValueError:
+        pass
+
+
+def test_parse_args_are_passed_along():
+    for option in all_arguments():
+        if option.cli_name != 'mask':
+            yield check_argument, option.cli_name, option.init_name, option.pass_value
+
+
+def test_parse_arg_types():
+    for option in ARGUMENT_SPEC_TYPED:
+        yield check_argument_type, option.cli_name, option.fail_value
+
+
+def test_check_duplicate_color_error():
+    color_mask_file = NamedTemporaryFile()
+    text_file = NamedTemporaryFile()
+
+    try:
+        cli.parse_args(['--color', 'red', '--colormask', color_mask_file.name, '--text', text_file.name])
+        raise AssertionError('parse_args(...) didn\'t raise')
+    except ValueError as e:
+        assert_true('specify either' in str(e), msg='expecting the correct error message, instead got: ' + str(e))
+
+
+def test_parse_args_defaults_to_random_color():
+    text = NamedTemporaryFile()
+
+    args = cli.parse_args(['--text', text.name])
+    assert_equal(args.color_func, wc.random_color_func)

wordcloud/TODO

-* command line interface
 * html export
 * good notebook interface
 * by default differnt color schemes
@@ -6,4 +5,3 @@
 * long functions?
 * redo examples
 * examples
-* filter one-letter words

wordcloud/__init__.py

-from .wordcloud import WordCloud, STOPWORDS, random_color_func
+from .wordcloud import WordCloud, STOPWORDS, random_color_func, get_single_color_func
 from .color_from_image import ImageColorGenerator
 
-__all__ = ['WordCloud', 'STOPWORDS', 'random_color_func', 'ImageColorGenerator']
+__all__ = ['WordCloud', 'STOPWORDS', 'random_color_func', 'get_single_color_func', 'ImageColorGenerator']

wordcloud/wordcloud.py

@@ -10,10 +10,12 @@ from random import Random
 import os
 import re
 import sys
+import colorsys
 import numpy as np
 from operator import itemgetter
 
 from PIL import Image
+from PIL import ImageColor
 from PIL import ImageDraw
 from PIL import ImageFont
 
@@ -77,6 +79,37 @@ def random_color_func(word=None, font_size=None, position=None,
         random_state = Random()
     return "hsl(%d, 80%%, 50%%)" % random_state.randint(0, 255)
 
+def get_single_color_func(color):
+    """Create a color function which returns a single hue and saturation with.
+    different values (HSV). Accepted values are color strings as usable by PIL/Pillow.
+
+    >>> color_func1 = get_single_color_func('deepskyblue')
+    >>> color_func2 = get_single_color_func('#00b4d2')
+    """
+    old_r, old_g, old_b = ImageColor.getrgb(color)
+    rgb_max = 255.
+    h, s, v = colorsys.rgb_to_hsv(old_r/rgb_max, old_g/rgb_max, old_b/rgb_max)
+    def single_color_func(word=None, font_size=None, position=None,
+                          orientation=None, font_path=None, random_state=None):
+        """Random color generation.
+
+        Additional coloring method. It picks a random value with hue and
+        saturation based on the color given to the generating function.
+
+        Parameters
+        ----------
+        word, font_size, position, orientation  : ignored.
+
+        random_state : random.Random object or None, (default=None)
+          If a random object is given, this is used for generating random numbers.
+
+        """
+        if random_state is None:
+            random_state = Random()
+        r, g, b = colorsys.hsv_to_rgb(h, s, random_state.uniform(0.2, 1))
+        return 'rgb({:.0f}, {:.0f}, {:.0f})'.format(r * rgb_max, g * rgb_max, b * rgb_max)
+    return single_color_func
+
 
 class WordCloud(object):
     """Word cloud object for generating and drawing.
@@ -350,7 +383,7 @@ class WordCloud(object):
         d = {}
         flags = (re.UNICODE if sys.version < '3' and type(text) is unicode
                  else 0)
-        for word in re.findall(r"\w[\w']*", text, flags=flags):
+        for word in re.findall(r"\w[\w']+", text, flags=flags):
             if word.isdigit():
                 continue
 

wordcloud/wordcloud_cli.py

+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+r"""Command-line tool to generate word clouds
+Usage::
+    $ cat word.txt | wordcloud_cli.py
+
+    $ wordcloud_cli.py --text=words.txt --stopwords=stopwords.txt
+"""
+import argparse
+import wordcloud as wc
+import numpy as np
+import sys
+import io
+from PIL import Image
+
+def main(args):
+    wordcloud = wc.WordCloud(stopwords=args.stopwords, mask=args.mask,
+        width=args.width, height=args.height, font_path=args.font_path,
+        margin=args.margin, relative_scaling=args.relative_scaling,
+        color_func=args.color_func, background_color=args.background_color).generate(args.text)
+    image = wordcloud.to_image()
+
+    b = io.BytesIO()
+    image.save(b, format='png')
+    with args.imagefile:
+        args.imagefile.write(b.getvalue())
+
+def parse_args(arguments):
+    prog = 'python wordcloud_cli.py'
+    description = ('A simple command line interface for wordcloud module.')
+    parser = argparse.ArgumentParser(description=description)
+    parser.add_argument('--text', metavar='file', type=argparse.FileType(), default='-',
+        help='specify file of words to build the word cloud (default: stdin)')
+    parser.add_argument('--stopwords', metavar='file', type=argparse.FileType(),
+        help='specify file of stopwords (containing one word per line) to remove from the given text after parsing')
+    parser.add_argument('--imagefile', metavar='file', type=argparse.FileType('w'), default='-',
+        help='file the completed PNG image should be written to (default: stdout)')
+    parser.add_argument('--fontfile', metavar='path', dest='font_path',
+        help='path to font file you wish to use (default: DroidSansMono)')
+    parser.add_argument('--mask', metavar='file', type=argparse.FileType(),
+        help='mask to use for the image form')
+    parser.add_argument('--colormask', metavar='file', type=argparse.FileType(),
+        help='color mask to use for image coloring')
+    parser.add_argument('--relative_scaling', type=float, default=0,
+        metavar='rs', help=' scaling of words by frequency (0 - 1)')
+    parser.add_argument('--margin', type=int, default=2,
+        metavar='width', help='spacing to leave around words')
+    parser.add_argument('--width', type=int, default=400,
+        metavar='width', help='define output image width')
+    parser.add_argument('--height', type=int, default=200,
+        metavar='height', help='define output image height')
+    parser.add_argument('--color', metavar='color',
+        help='use given color as coloring for the image - accepts any value from PIL.ImageColor.getcolor')
+    parser.add_argument('--background', metavar='color', default='black', type=str, dest='background_color',
+        help='use given color as background color for the image - accepts any value from PIL.ImageColor.getcolor')
+    args = parser.parse_args(arguments)
+
+    if args.colormask and args.color:
+        raise ValueError('specify either a color mask or a color function')
+
+    with args.text:
+        args.text = args.text.read()
+
+    if args.stopwords:
+        with args.stopwords:
+            args.stopwords = set(map(str.strip, args.stopwords.readlines()))
+
+    if args.mask:
+        args.mask = np.array(Image.open(args.mask))
+
+    color_func = wc.random_color_func
+    if args.colormask:
+        image = np.array(Image.open(args.colormask))
+        color_func = wc.ImageColorGenerator(image)
+
+    if args.color:
+        color_func = wc.get_single_color_func(args.color)
+
+    args.color_func = color_func
+    return args
+
+if __name__ == '__main__':
+    main(parse_args(sys.argv[1:]))