Merge pull request amueller#243 from cjmay/241-optional-collocations

Make collocations in word cloud optional in CLI.
Kitter · Mar 21, 2017 · c9402a1 · c9402a1
2 parents 2b86894 + 82da2d6
commit c9402a1
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/test/test_wordcloud_cli.py b/test/test_wordcloud_cli.py
@@ -17,7 +17,10 @@
     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')
+    ArgOption(cli_name='relative_scaling', init_name='relative_scaling', pass_value=1, fail_value='c'),
+]
+ARGUMENT_SPEC_UNARY = [
+    ArgOption(cli_name='no_collocations', init_name='collocations', pass_value=True, fail_value=1)
 ]
 ARGUMENT_SPEC_REMAINING = [
     ArgOption(cli_name='stopwords', init_name='stopwords', pass_value=temp.name, fail_value=None),
@@ -26,10 +29,12 @@
     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)
 ]
+ARGUMENT_CLI_NAMES_UNARY = [arg_opt.cli_name for arg_opt in ARGUMENT_SPEC_UNARY]
 
 def all_arguments():
     arguments = []
     arguments.extend(ARGUMENT_SPEC_TYPED)
+    arguments.extend(ARGUMENT_SPEC_UNARY)
     arguments.extend(ARGUMENT_SPEC_REMAINING)
     return arguments
 
@@ -56,6 +61,13 @@ def check_argument(name, result_name, value):
     assert_in(result_name, vars(args))
 
 
+def check_argument_unary(name, result_name):
+    text = NamedTemporaryFile()
+
+    args = cli.parse_args(['--text', text.name, '--' + name])
+    assert_in(result_name, vars(args))
+
+
 def check_argument_type(name, value):
     text = NamedTemporaryFile()
 
@@ -71,7 +83,9 @@ def check_argument_type(name, value):
 
 def test_parse_args_are_passed_along():
     for option in all_arguments():
-        if option.cli_name != 'mask':
+        if option.cli_name in ARGUMENT_CLI_NAMES_UNARY:
+            yield check_argument_unary, option.cli_name, option.init_name
+        elif option.cli_name != 'mask':
             yield check_argument, option.cli_name, option.init_name, option.pass_value
 
 

diff --git a/wordcloud/wordcloud_cli.py b/wordcloud/wordcloud_cli.py
@@ -16,7 +16,8 @@ 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)
+        color_func=args.color_func, background_color=args.background_color,
+        collocations=args.collocations).generate(args.text)
     image = wordcloud.to_image()
 
     with args.imagefile:
@@ -51,6 +52,8 @@ def parse_args(arguments):
         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')
+    parser.add_argument('--no_collocations', action='store_true',
+        help='do not add collocations (bigrams) to word cloud (default: add unigrams and bigrams)')
     args = parser.parse_args(arguments)
 
     if args.colormask and args.color:
@@ -74,6 +77,8 @@ def parse_args(arguments):
     if args.color:
         color_func = wc.get_single_color_func(args.color)
 
+    args.collocations = not args.no_collocations
+
     args.color_func = color_func
     return args