"""
 Playing around with Vigenere ciphers.

 (See http://en.wikipedia.org/wiki/Vigenere_cipher)

 >>> hello = Letters('hello')
 >>> print hello
 HELLO
 >>> print hello.encode('ab')
 HFLMO
 >>> print hello.encode('ab').decode('ab')
 HELLO

 >>> print Letters('Hows it going? Yeah, I thought so.').encode('strange')
 ZHNSVZKGBEGLKEZBKHBAKZMJO
 >>> print Letters('ZHNSVZKGBEGLKEZBKHBAKZMJO').decode('strange')
 HOWSITGOINGYEAHITHOUGHTSO

 Jim Mahoney | March 2012 | MIT License

"""

def letters_only(chars):
    """ Return only letters, removing numbers, spaces, punctuation.
        Or if chars is a list of Letter objects, just return it.
        >>> letters_only("One 2 Three.")
        'OneThree'
    """
    import re
    if len(chars) > 0 and isinstance(chars[0], Letter):
        return chars
    else:
        return ''.join(re.findall("[a-zA-Z]", chars))

class Letter:
    """ A-Z with wrap-around addition """
    def __init__(self, letter):
        self.letter = letter.upper()
    def __str__(self):
        return self.letter
    def upper(self):
        # A trick to make Letter(Letter('a')) same as Letter('a')
        return self.letter
    def value(self):
        return ord(self.letter) - ord("A")
    def addsub(self, other, sign):
        return Letter(chr(
                (self.value() + sign * other.value()) % 26 + ord("A")))
    def __add__(self, other):
        return self.addsub(other, 1)
    def __sub__(self, other):
        return self.addsub(other, -1)

class Letters:
    def __init__(self, chars=''):
        self.letters = [Letter(c) for c in letters_only(chars)]
    def __len__(self):
        return len(self.letters)
    def __str__(self):
        return ''.join([str(l) for l in self.letters])
    def __getitem__(self, n):
        if type(n)==slice:
            (start, stop, step) = (n.start or 0, n.stop, n.step or 1)
            return Letters([self.letters[i] for i in range(start, stop, step)])
        else:
            return self.letters[n]
    def __iter__(self):
        self.n_iter = -1
        return self
    def next(self):
        self.n_iter += 1
        if self.n_iter >= len(self):
            raise StopIteration
        else:
            return self[self.n_iter]
    def zipkey(self, key):
        """ Return [(letter0, key0), (letter1, key1), ...] """
        # So if str(self) is "HELLO" and key is "ab",
        # and A means Letter("A"), then what is returned is
        # [(H, 'a'), (E, 'b'), (L, 'a'), (L, 'b'), (O, 'a')]
        return zip(self.letters, key * (len(self)/len(key) + 1))
    def encode(self, key):
        return Letters([l + Letter(c) for (l, c) in self.zipkey(key)])
    def decode(self, key):
        return Letters([l - Letter(c) for (l, c) in self.zipkey(key)])

if __name__ == "__main__":
    import doctest
    doctest.testmod()