script for saving pdb chains in single file

scripts/chain_extract.py

@@ -10,21 +10,15 @@
 
 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())
-
+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