option to shuffle

1 files changed, 14 insertions, 7 deletions

diff --git a/cli/commands/bridge/words.py b/cli/commands/bridge/words.py
index 4bc0982..27d0bd4 100644
--- a/cli/commands/bridge/words.py
+++ b/cli/commands/bridge/words.py
@@ -30,13 +30,15 @@ from app.thesaurus.api import Thesaurus
   help='Minimum depth of matches')
 @click.option('-sh', '--use_shortest_path', 'opt_use_shortest_path', is_flag=True,
   help='Use shortest path between words')
+@click.option('--shuffle/--no_shuffle', 'opt_shuffle', default=True,
+  help='Whether to shuffle the input')
 @click.pass_context
-def cli(ctx, opt_word_a, opt_word_b, opt_include_oe, opt_include_slang, opt_include_scots, opt_words_per_step, opt_categories_per_word, opt_min_depth, opt_use_shortest_path):
+def cli(ctx, opt_word_a, opt_word_b, opt_include_oe, opt_include_slang, opt_include_scots, opt_words_per_step, opt_categories_per_word, opt_min_depth, opt_use_shortest_path, opt_shuffle):
   """
   Find connections between two words
   """
   thesaurus = Thesaurus()
-  solver = TreeSolver(thesaurus, opt_word_a, opt_word_b, opt_include_oe, opt_include_slang, opt_include_scots, opt_words_per_step, opt_categories_per_word, opt_min_depth, opt_use_shortest_path)
+  solver = TreeSolver(thesaurus, opt_word_a, opt_word_b, opt_include_oe, opt_include_slang, opt_include_scots, opt_words_per_step, opt_categories_per_word, opt_min_depth, opt_use_shortest_path, opt_shuffle)
   print(f"Starting word: {opt_word_a}")
   print(f"Ending word: {opt_word_b}")
 
@@ -58,7 +60,7 @@ def cli(ctx, opt_word_a, opt_word_b, opt_include_oe, opt_include_slang, opt_incl
     # print(solver.skips)
 
 class TreeSolver:
-  def __init__(self, thesaurus, word_a, word_b, include_oe, include_slang, include_scots, words_per_step, categories_per_word, min_depth, use_shortest_path):
+  def __init__(self, thesaurus, word_a, word_b, include_oe, include_slang, include_scots, words_per_step, categories_per_word, min_depth, use_shortest_path, shuffle):
     self.thesaurus = thesaurus
     self.word_a = word_a
     self.word_b = word_b
@@ -71,6 +73,7 @@ class TreeSolver:
     self.min_depth = min_depth
     self.use_shortest_path = use_shortest_path
     self.max_dist = 0
+    self.shuffle = shuffle
     self.reset()
 
   def reset(self):
@@ -85,7 +88,8 @@ class TreeSolver:
       words = words[:self.words_per_step]
     for word in tqdm(words):
       categories = self.thesaurus.search(word)['categories']
-      random.shuffle(categories)
+      if self.shuffle:
+        random.shuffle(categories)
       count = 0
       for category in categories:
         if count > self.categories_per_word:
@@ -103,7 +107,8 @@ class TreeSolver:
         if len(add_to_queue):
           next_queue += add_to_queue
           count += 1
-    random.shuffle(next_queue)
+    if self.shuffle:
+      random.shuffle(next_queue)
     return next_queue
 
   def process_category(self, catid, tree, target):
@@ -182,7 +187,8 @@ class TreeSolver:
       match = None
       if self.is_integer(word):
         category_words = self.thesaurus.category(word)['words']
-        random.shuffle(category_words)
+        if self.shuffle:
+          random.shuffle(category_words)
         for category_word in category_words:
           cat_word = self.fix_word(category_word['word'])
           if cat_word != word and cat_word in tree and self.can_descend(tree, cat_word, word):
@@ -191,7 +197,8 @@ class TreeSolver:
             break
       else:
         categories = self.thesaurus.search(word)['categories']
-        random.shuffle(categories)
+        if self.shuffle:
+          random.shuffle(categories)
         for category in categories:
           catid = category['catid']
           if catid != word and catid in tree and self.can_descend(tree, catid, word):