Nucleotide check.py

def part_seq(seq, size):
    """Split a sequence in parts of length size and return
    a list of lists where the latter have length size.
    """
    assert len(seq)>0
    parts = []
    for i in range(0,len(seq), size):
        parts.append(seq[i:i+size])
    return parts

def is_valid(codon):
    return check1(codon) or check2(codon)

def check1(codon):
    """Check 1: All characters the same"""
    for c in codon:
        if c != codon[0]:
            return False
    return True

def check2(codon):
    """Check 2: All characters are enumeration"""
    def make_ordering(orderstr):
        order = {}
        for i, c in enumerate(orderstr):
            order[c] = i
        return order
    order = make_ordering('ABCDHGKMR')
    first = codon[0]
    for i, c in enumerate(codon):
        if order[c]-order[first] != i:
            return False
    return True


def check_nucleotide_sequence(s):  
    assert len(s) == 12
    codons = part_seq(s, 3) #Split seq in parts of length 3
    for codon in codons:
        if not is_valid(codon):
            return False
    return True

if __name__ == '__main__':
    """Tests"""
    assert check_nucleotide_sequence("AAABBBCCCDDD")
    assert check_nucleotide_sequence("ABCAAABBBBCD")
    assert check2("ABC")
    assert check1("AAA")
    assert is_valid("AAA")
    assert is_valid("ABC")
    print('ALL OK')