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')))

def set_numbering(selection='all', firstname='N', start=1):
	'''
DESCRIPTION

    Set residue numbering based on current order of atoms.
	'''
	from pymol import stored
	stored.resi = int(start)-1
	def guess_resi(name):
		if name == firstname:
			stored.resi += 1
		return stored.resi
	cmd.alter(selection, 'resi = guess_resi(name)', space=locals())

cmd.extend('set_sequence', set_sequence)
cmd.extend('set_numbering', set_numbering)