anagrams.py

import unittest

from collections import defaultdict


class Anagrams(object):
    def __init__(self):
        self.words = open('words.txt').readlines()
        self._load_anagrams()

    def _load_anagrams(self):
        self.anagrams = defaultdict(list)
        for word in self.words:
            word = word.strip()
            key = "".join(sorted(word))
            self.anagrams[key].append(word)

    def get_anagrams(self, word):
        return self.anagrams["".join(sorted(word))]


class TestAnagrams(unittest.TestCase):
    def test_anagrams(self):
        anagrams = Anagrams()
        self.assertEquals(anagrams.get_anagrams('plates'), ['palest', 'pastel', 'petals', 'plates', 'staple'])
        self.assertEquals(anagrams.get_anagrams('eat'), ['ate', 'eat', 'tea'])

    def test_anagrams_non_existant(self):
        anagrams = Anagrams()
        self.assertEquals(anagrams.get_anagrams('nonexistent'), [])


if __name__ == '__main__':
    unittest.main()