#!/usr/bin/python3

#=======================================================================
# TASK: select specified chains from the pdb & save a cropped PDB with
# the selected chains. Useful for dimer, etc modelling.

# link for saving each chain as a separate file
# https://stackoverflow.com/questions/11685716/how-to-extract-chains-from-a-pdb-file
#=======================================================================

from Bio.PDB import PDBParser, PDBIO, Select

io = PDBIO()
pdb = PDBParser().get_structure("3byw", "3byw.pdb")


# Select() Method to return True for every chain in 'chains'
class ChainSelect(Select):
   def accept_chain(self, chain):
       #print dir(chain)
       if chain.id in chains:
           return 1
       else:
           return 0

if __name__ == '__main__':
    chains = ['G', 'H'] # specify selected chains
    io.set_structure(pdb)
    io.save(pdb.get_id() + "_crop.pdb", ChainSelect())