renumber_pdb_str_by_length_list.py

def renum_pdb_str(pdb_str, Ls=None, renum=True, offset=1):
    from string import ascii_uppercase, ascii_lowercase
    assert len(Ls) == 2
    Ls = [int(i) for i in Ls]
    alphabet_list = list(ascii_uppercase+ascii_lowercase)
    if Ls is not None:
        L_init = 0
        new_chain = {}
        for L,c in zip(Ls, alphabet_list):
            new_chain.update({i:c for i in range(L_init,L_init+L)})
            L_init += L  

    n,num,pdb_out = 0,offset,[]
    resnum_ = None
    chain_ = None
    new_chain_ = new_chain[0]
    for line in pdb_str.split("\n"):
        if line[:4] == "ATOM":
            chain = line[21:22]
            resnum = int(line[22:22+5])
            if resnum_ is None: resnum_ = resnum
            if chain_ is None: chain_ = chain
            if resnum != resnum_ or chain != chain_:
                num += (resnum - resnum_)  
                n += 1
                resnum_,chain_ = resnum,chain
            if Ls is not None:
                if new_chain[n] != new_chain_:
                    num = offset
                    new_chain_ = new_chain[n]
            N = num if renum else resnum
            if Ls is None: pdb_out.append("%s%4i%s" % (line[:22],N,line[26:]))
            else: pdb_out.append("%s%s%4i%s" % (line[:21],new_chain[n],N,line[26:]))        
            
    pdb_str = "\n".join(pdb_out)
    return pdb_str