1import argparse
2import os
3import subprocess
4from collections import namedtuple
5
6import wordcloud as wc
7from wordcloud import wordcloud_cli as cli
8
9from mock import patch
10import pytest
11
12import matplotlib
13matplotlib.use('Agg')
14
15
16class PassFile(object):
17    pass
18
19
20ArgOption = namedtuple('ArgOption', ['cli_name', 'init_name', 'pass_value', 'fail_value'])
21ARGUMENT_SPEC_TYPED = [
22    ArgOption(cli_name='width', init_name='width', pass_value=13, fail_value=1.),
23    ArgOption(cli_name='height', init_name='height', pass_value=15, fail_value=1.),
24    ArgOption(cli_name='margin', init_name='margin', pass_value=17, fail_value=1.),
25    ArgOption(cli_name='relative_scaling', init_name='relative_scaling', pass_value=1, fail_value='c'),
26]
27ARGUMENT_SPEC_UNARY = [
28    ArgOption(cli_name='no_collocations', init_name='collocations', pass_value=True, fail_value=1)
29]
30ARGUMENT_SPEC_REMAINING = [
31    ArgOption(cli_name='stopwords', init_name='stopwords', pass_value=PassFile(), fail_value=None),
32    ArgOption(cli_name='regexp', init_name='regexp', pass_value=r'\w{2,}', fail_value=r'12('),
33    ArgOption(cli_name='mask', init_name='mask', pass_value=PassFile(), fail_value=None),
34    ArgOption(cli_name='fontfile', init_name='font_path', pass_value=PassFile(), fail_value=None),
35    ArgOption(cli_name='color', init_name='color_func', pass_value='red', fail_value=None),
36    ArgOption(cli_name='background', init_name='background_color', pass_value='grey', fail_value=None),
37    ArgOption(cli_name='contour_color', init_name='contour_color', pass_value='grey', fail_value=None),
38    ArgOption(cli_name='contour_width', init_name='contour_width', pass_value=0.5, fail_value='blue'),
39]
40ARGUMENT_CLI_NAMES_UNARY = [arg_opt.cli_name for arg_opt in ARGUMENT_SPEC_UNARY]
41
42
43def all_arguments():
44    arguments = []
45    arguments.extend(ARGUMENT_SPEC_TYPED)
46    arguments.extend(ARGUMENT_SPEC_UNARY)
47    arguments.extend(ARGUMENT_SPEC_REMAINING)
48    return arguments
49
50
51def test_main_passes_arguments_through(tmpdir):
52
53    image_filepath = str(tmpdir.join('word_cloud.png'))
54
55    args = argparse.Namespace()
56    for option in all_arguments():
57        setattr(args, option.init_name, option.pass_value)
58
59    text = 'some long text'
60    image_file = open(image_filepath, 'w')
61    with patch('wordcloud.wordcloud_cli.wc.WordCloud', autospec=True) as mock_word_cloud:
62        cli.main(vars(args), text, image_file)
63
64    posargs, kwargs = mock_word_cloud.call_args
65    for option in all_arguments():
66        assert option.init_name in kwargs
67
68
69def check_argument(text_filepath, name, result_name, value):
70    args, text, image_file = cli.parse_args(['--text', text_filepath, '--' + name, str(value)])
71    assert result_name in args
72
73
74def check_argument_unary(text_filepath, name, result_name):
75    args, text, image_file = cli.parse_args(['--text', text_filepath, '--' + name])
76    assert result_name in args
77
78
79def check_argument_type(text_filepath, name, value):
80    with pytest.raises(
81            (SystemExit, ValueError),
82            message='argument "{}" was accepted even though the type did not match'.format(name)
83    ):
84        args, text, image_file = cli.parse_args(['--text', text_filepath, '--' + name, str(value)])
85
86
87@pytest.mark.parametrize("option", all_arguments())
88def test_parse_args_are_passed_along(option, tmpdir, tmp_text_file):
89    if option.cli_name in ARGUMENT_CLI_NAMES_UNARY:
90        check_argument_unary(str(tmp_text_file), option.cli_name, option.init_name)
91    elif option.cli_name != 'mask':
92        pass_value = option.pass_value
93        if isinstance(option.pass_value, PassFile):
94            input_file = tmpdir.join("%s_file" % option.cli_name)
95            input_file.write(b"")
96            pass_value = str(input_file)
97        check_argument(str(tmp_text_file), option.cli_name, option.init_name, pass_value)
98
99
100@pytest.mark.parametrize("option", ARGUMENT_SPEC_TYPED)
101def test_parse_arg_types(option, tmp_text_file):
102    check_argument_type(str(tmp_text_file), option.cli_name, option.fail_value)
103
104
105def test_check_duplicate_color_error(tmpdir, tmp_text_file):
106    color_mask_file = tmpdir.join("input_color_mask.png")
107    color_mask_file.write(b"")
108
109    with pytest.raises(ValueError, match=r'.*specify either.*'):
110        cli.parse_args(['--color', 'red', '--colormask', str(color_mask_file), '--text', str(tmp_text_file)])
111
112
113def test_parse_args_defaults_to_random_color(tmp_text_file):
114    args, text, image_file = cli.parse_args(['--text', str(tmp_text_file)])
115    assert args['color_func'] == wc.random_color_func
116
117
118def test_unicode_text_file():
119    unicode_file = os.path.join(os.path.dirname(__file__), "unicode_text.txt")
120    args, text, image_file = cli.parse_args(['--text', unicode_file])
121    assert len(text) == 16
122
123
124def test_cli_writes_image(tmpdir, tmp_text_file):
125    # ensure writing works with all python versions
126    tmp_image_file = tmpdir.join("word_cloud.png")
127
128    tmp_text_file.write(b'some text')
129
130    args, text, image_file = cli.parse_args(['--text', str(tmp_text_file), '--imagefile', str(tmp_image_file)])
131    cli.main(args, text, image_file)
132
133    # expecting image to be written
134    assert tmp_image_file.size() > 0
135
136
137def test_cli_regexp(tmp_text_file):
138    cli.parse_args(['--regexp', r"\w[\w']+", '--text', str(tmp_text_file)])
139
140
141def test_cli_regexp_invalid(tmp_text_file, capsys):
142    with pytest.raises(SystemExit):
143        cli.parse_args(['--regexp', r"invalid[", '--text', str(tmp_text_file)])
144
145    _, err = capsys.readouterr()
146    assert "Invalid regular expression" in err
147
148
149@pytest.mark.parametrize("command,expected_output, expected_exit_code", [
150    ("wordcloud_cli --help", "usage: wordcloud_cli", 0),
151    ("python -m wordcloud --help", "usage: __main__", 0),
152    ("python %s/../wordcloud/wordcloud_cli.py --help" % os.path.dirname(__file__), "To execute the CLI", 1),
153])
154def test_cli_as_executable(command, expected_output, expected_exit_code, tmpdir, capfd, no_cover_compat):
155
156    ret_code = 0
157    try:
158        subprocess.check_call(
159            command,
160            shell=True,
161            cwd=str(tmpdir)
162        )
163    except subprocess.CalledProcessError as excinfo:
164        ret_code = excinfo.returncode
165
166    out, err = capfd.readouterr()
167    assert expected_output in out if ret_code == 0 else err
168
169    assert ret_code == expected_exit_code
170