#!/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


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