1 files changed, 24 insertions, 22 deletions

diff --git a/labb2/kod.py b/labb2/kod.py
index 752ba43..378225f 100644
--- a/labb2/kod.py
+++ b/labb2/kod.py
@@ -1,23 +1,25 @@
 from lib import minimum_edit_distance
 
 
-words = {}
-for line in open("alphabetical.csv", "r").readlines():
-    word, freq = line.split(",")
-    words[word] = int(freq)
-
-
-def main():
+def load():
+    words = {}
+    for line in open("alphabetical.csv", "r").readlines():
+        word, freq = line.split(",")
+        words[word] = int(freq)
+    return words
+
+def main(words):
+    """Repeatedly ask the user for word and autocomplete/-correct it."""
     word = ""
     while word != "q":
         word = input("Type word: ").lower()
-        print("Autocompletion finished: ", autocomplete(word))
-        print("Sorted autocompletion: ", autocomplete_best(word))
-        print("Best three: ", autocomplete_best(word, 3))
-        print("Autocorrect: ", autocorrect(word))
+        print("Autocompletion finished: ", autocomplete(words, word))
+        print("Sorted autocompletion: ", autocomplete_best(words, word))
+        print("Best three: ", autocomplete_best(words, word, 3))
+        print("Autocorrect: ", autocorrect(words, word))
 
 
-def autocomplete(search_word):
+def autocomplete(words, search_word):
     """Return autocomplete suggestions."""
     for word in words.keys():
         if word.startswith(search_word):
@@ -25,21 +27,22 @@ def autocomplete(search_word):
     return None
 
 
-def autocomplete_best(search_word, amount=None):
+def autocomplete_best(words, search_word, amount=None):
     """Return the /amount/ most common autocompletions."""
     matching_words = {word: freq for word, freq in words.items()
                       if word.startswith(search_word)}
-    matching_words_sorted = {word: freq for word, freq in
-                             sorted(matching_words.items(),
-                                    key=lambda item: item[1],
-                                    reverse=True)}
+    matching_words_sorted = sorted(matching_words.items(),
+                                   key=lambda item: item[1],  # sort by freq
+                                   reverse=True)  # sort most frequent first
+    matching_words_sorted_words = [word for word, freq in matching_words_sorted]
+
     if amount:
-        return ", ".join(list(matching_words_sorted.keys())[:amount])
+        return ", ".join(list(matching_words_sorted_words)[:amount])
     else:
-        return ", ".join(list(matching_words_sorted.keys()))
+        return ", ".join(list(matching_words_sorted_words))
 
 
-def autocorrect(search_word):
+def autocorrect(words, search_word):
     """Return the word with the smallest Levhenstein distance"""
     best = None
     for word, _ in words.items():
@@ -50,5 +53,4 @@ def autocorrect(search_word):
             best = (edit_distance, word)
     return best[1]
 
-
-main()
+main(load())