diff --git a/src/wormhole/_wordlist.py b/src/wormhole/_wordlist.py index ac48e78..cf801c5 100644 --- a/src/wormhole/_wordlist.py +++ b/src/wormhole/_wordlist.py @@ -1,4 +1,4 @@ -from __future__ import unicode_literals +from __future__ import unicode_literals, print_function import os from zope.interface import implementer from ._interfaces import IWordlist @@ -160,9 +160,10 @@ for k,both_words in raw_words.items(): @implementer(IWordlist) class PGPWordList(object): - def get_completions(self, prefix): + def get_completions(self, prefix, num_words=2): # start with the odd words - if prefix.count("-") % 2 == 0: + count = prefix.count("-") + if count % 2 == 0: words = odd_words_lowercase else: words = even_words_lowercase @@ -171,7 +172,11 @@ class PGPWordList(object): completions = set() for word in words: if word.startswith(last_partial_word): - completions.add(word[lp:]) + suffix = word[lp:] + # append a hyphen if we expect more words + if count+1 < num_words: + suffix += "-" + completions.add(suffix) return completions def choose_words(self, length): diff --git a/src/wormhole/test/test_wordlist.py b/src/wormhole/test/test_wordlist.py index 56cac4b..c86fbdb 100644 --- a/src/wormhole/test/test_wordlist.py +++ b/src/wormhole/test/test_wordlist.py @@ -7,11 +7,16 @@ class Completions(unittest.TestCase): def test_completions(self): wl = PGPWordList() gc = wl.get_completions - self.assertEqual(gc("ar"), {"mistice", "ticle"}) - self.assertEqual(gc("armis"), {"tice"}) - self.assertEqual(gc("armistice-ba"), + self.assertEqual(gc("ar", 2), {"mistice-", "ticle-"}) + self.assertEqual(gc("armis", 2), {"tice-"}) + self.assertEqual(gc("armistice", 2), {"-"}) + self.assertEqual(gc("armistice-ba", 2), {"boon", "ckfield", "ckward", "njo"}) - self.assertEqual(gc("armistice-baboon"), {""}) + self.assertEqual(gc("armistice-ba", 3), + {"boon-", "ckfield-", "ckward-", "njo-"}) + self.assertEqual(gc("armistice-baboon", 2), {""}) + self.assertEqual(gc("armistice-baboon", 3), {"-"}) + self.assertEqual(gc("armistice-baboon", 4), {"-"}) class Choose(unittest.TestCase): def test_choose_words(self):