from pymol import cmd

def set_sequence(sequence, selection='all', start=1):
    '''
DESCRIPTION

    Alters the residue names according to given sequence

ARGUMENTS

    sequence = string: amino acid sequence in one-letter code

    selection = string: atom selection {default: all}

    start = int: residue number to start from {default: 1}
    '''
    import re
    from pymol.exporting import _resn_to_aa as one_letter
    three_letter = dict((v,k) for (k,v) in one_letter.iteritems())
    sequence = re.sub(r'\s+', '', sequence)
    start = int(start)
    for i, aa in enumerate(sequence):
        cmd.alter('(%s) and resi %d' % (selection, i+start),
                'resn=' + repr(three_letter.get(aa.upper(), 'UNK')))

cmd.extend('set_sequence', set_sequence)
