From 43820d8fb44c6c7deba216abdc81117bc3bc250f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Rodrigues?= <j.p.g.l.m.rodrigues@gmail.com>
Date: Mon, 20 Jun 2016 00:24:32 -0700
Subject: [PATCH] Added support for multiple chains/segments in selection
 scripts.

---
 pdb_selchain.py | 10 +++++-----
 pdb_selseg.py   | 10 +++++-----
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/pdb_selchain.py b/pdb_selchain.py
index 11e0a9aa..28b3e02b 100755
--- a/pdb_selchain.py
+++ b/pdb_selchain.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 
 """
-Extracts a chain from a PDB file.
+Extracts one or more chains from a PDB file.
 
 usage: python pdb_selchain.py -<chain> <pdb file>
 example: python pdb_selchain.py -A 1CTF.pdb
@@ -37,7 +37,7 @@ def check_input(args):
             sys.exit(1)
     elif len(args) == 1:
         # Chain & Pipe _or_ file & no chain
-        if re.match('\-[A-Za-z0-9]', args[0]):
+        if re.match('\-[A-Za-z0-9]+', args[0]):
             chain = args[0][1:]
             if not sys.stdin.isatty():
                 pdbfh = sys.stdin
@@ -53,7 +53,7 @@ def check_input(args):
             chain = ' '
     elif len(args) == 2:
         # Chain & File
-        if not re.match('\-[A-Za-z0-9]', args[0]):
+        if not re.match('\-[A-Za-z0-9]+', args[0]):
             sys.stderr.write('Invalid chain ID: ' + args[0] + '\n')
             sys.stderr.write(USAGE)
             sys.exit(1)
@@ -74,10 +74,10 @@ def _select_chain(fhandle, chain_id):
 
     coord_re = re.compile('^(ATOM|HETATM)')
     fhandle = fhandle
-    chain_id = chain_id
+    chain_id = set(chain_id)
 
     for line in fhandle:
-        if coord_re.match(line) and line[21] == chain_id:
+        if coord_re.match(line) and line[21] in chain_id:
             yield line
 
 if __name__ == '__main__':
diff --git a/pdb_selseg.py b/pdb_selseg.py
index c5908c62..c6bbcdf1 100755
--- a/pdb_selseg.py
+++ b/pdb_selseg.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 
 """
-Extracts a segment from a PDB file.
+Extracts one or more segments from a PDB file.
 
 usage: python pdb_selseg.py -<segid> <pdb file>
 example: python pdb_selseg.py -A 1CTF.pdb
@@ -37,7 +37,7 @@ def check_input(args):
             sys.exit(1)
     elif len(args) == 1:
         # Segid & Pipe _or_ file & no Segid
-        if re.match('\-[A-Za-z0-9]', args[0]):
+        if re.match('\-[A-Za-z0-9]+', args[0]):
             seg = args[0][1:]
             if not sys.stdin.isatty():
                 pdbfh = sys.stdin
@@ -53,7 +53,7 @@ def check_input(args):
             seg = ' '
     elif len(args) == 2:
         # Segid & File
-        if not re.match('\-[A-Za-z0-9]', args[0]):
+        if not re.match('\-[A-Za-z0-9]+', args[0]):
             sys.stderr.write('Invalid segment ID: ' + args[0] + '\n')
             sys.stderr.write(USAGE)
             sys.exit(1)
@@ -74,10 +74,10 @@ def _select_seg(fhandle, seg_id):
 
     coord_re = re.compile('^(ATOM|HETATM)')
     fhandle = fhandle
-    seg_id = seg_id
+    seg_id = set(seg_id)
 
     for line in fhandle:
-        if coord_re.match(line) and line[72:76].strip() == seg_id:
+        if coord_re.match(line) and line[72:76].strip() in seg_id:
             yield line
 
 if __name__ == '__main__':