Merge pull request #109

Adjust fhandle names to files and iterables.

pdbtools/__init__.py

@@ -72,6 +72,7 @@
 >>> help(MODULE.run)
 """
 
+
 __all__ = [
     'pdb_b',
     'pdb_chainbows',

pdbtools/pdb_splitchain.py

@@ -34,6 +34,7 @@
 import os
 import sys
 
+
 __author__ = "Joao Rodrigues"
 __email__ = "[email protected]"
 
@@ -71,7 +72,7 @@ def check_input(args):
     return fh
 
 
-def run(fhandle):
+def run(fhandle, outname=None):
     """
     Split the PDB into its different chains.
 
@@ -81,9 +82,21 @@ def run(fhandle):
     Parameters
     ----------
     fhandle : an iterable giving the PDB file line-by-line
+
+    outname : str
+        The base name of the output files. If None is given, tries to
+        extract a name from the `.name` attribute of `fhandler`. If
+        `fhandler` has no attribute name, assigns `splitchains`.
     """
-    fname_root = fhandle.name[:-4] if fhandle.name != '<stdin>' else 'output'
-    basename = os.path.basename(fname_root)
+    _defname = 'splitchains'
+    if outname is None:
+        try:
+            fn = fhandle.name
+            outname = fn[:-4] if fn != '<stdin>' else _defname
+        except AttributeError:
+            outname = _defname
+
+    basename = os.path.basename(outname)
 
     chain_data = {}  # {chain_id: lines}
 

pdbtools/pdb_splitmodel.py

@@ -34,6 +34,7 @@
 import os
 import sys
 
+
 __author__ = "Joao Rodrigues"
 __email__ = "[email protected]"
 
@@ -71,7 +72,7 @@ def check_input(args):
     return fh
 
 
-def run(fhandle):
+def run(fhandle, outname=None):
     """
     Split PDB into MODELS.
 
@@ -81,9 +82,21 @@ def run(fhandle):
     Parameters
     ----------
     fhandle : a line-by-line iterator of the original PDB file.
+
+    outname : str
+        The base name of the output files. If None is given, tries to
+        extract a name from the `.name` attribute of `fhandler`. If
+        `fhandler` has no attribute name, assigns `splitmodels`.
     """
-    fname_root = fhandle.name[:-4] if fhandle.name != '<stdin>' else 'pdbfile'
-    basename = os.path.basename(fname_root)
+    _defname = 'splitmodels'
+    if outname is None:
+        try:
+            fn = fhandle.name
+            outname = fn[:-4] if fn != '<stdin>' else _defname
+        except AttributeError:
+            outname = _defname
+
+    basename = os.path.basename(outname)
 
     model_lines = []
     records = ('ATOM', 'HETATM', 'ANISOU', 'TER')

pdbtools/pdb_splitseg.py

@@ -34,6 +34,7 @@
 import os
 import sys
 
+
 __author__ = "Joao Rodrigues"
 __email__ = "[email protected]"
 
@@ -71,7 +72,7 @@ def check_input(args):
     return fh
 
 
-def run(fhandle):
+def run(fhandle, outname=None):
     """
     Split PDB into segments.
 
@@ -81,9 +82,21 @@ def run(fhandle):
     Parameters
     ----------
     fhandle : a line-by-line iterator of the original PDB file.
+
+    outname : str
+        The base name of the output files. If None is given, tries to
+        extract a name from the `.name` attribute of `fhandler`. If
+        `fhandler` has no attribute name, assigns `splitsegs`.
     """
-    fname_root = fhandle.name[:-4] if fhandle.name != '<stdin>' else 'output'
-    basename = os.path.basename(fname_root)
+    _defname = 'splitsegs'
+    if outname is None:
+        try:
+            fn = fhandle.name
+            outname = fn[:-4] if fn != '<stdin>' else _defname
+        except AttributeError:
+            outname = _defname
+
+    basename = os.path.basename(outname)
 
     segment_data = {}  # {segment_id: lines}
 

pdbtools/pdb_tocif.py

@@ -36,6 +36,7 @@
 import os
 import sys
 
+
 __author__ = "Joao Rodrigues"
 __email__ = "[email protected]"
 
@@ -80,7 +81,7 @@ def pad_line(line):
     return line[:81]  # 80 + newline character
 
 
-def run(fhandle):
+def run(fhandle, outname=None):
     """
     Convert a structure in PDB format to mmCIF format.
 
@@ -90,6 +91,11 @@ def run(fhandle):
     ----------
     fhandle : an iterable giving the PDB file line-by-line.
 
+    outname : str
+        The base name of the output files. If None is given, tries to
+        extract a name from the `.name` attribute of `fhandler`. If
+        `fhandler` has no attribute name, assigns `cell`.
+
     Yields
     ------
     str (line-by-line)
@@ -106,10 +112,17 @@ def run(fhandle):
     yield '#\n'
 
     # Headers
-    fname, _ = os.path.splitext(os.path.basename(fhandle.name))
-    if fname == '<stdin>':
-        fname = 'cell'
-    yield 'data_{}\n'.format(fname)
+    _defname = 'cell'
+    if outname is None:
+        try:
+            fn = fhandle.name
+            outname = fn[:-4] if fn != '<stdin>' else _defname
+        except AttributeError:
+            outname = _defname
+
+    fname_root = os.path.basename(outname)
+
+    yield 'data_{}\n'.format(fname_root)
 
     yield '#\n'
     yield 'loop_\n'

tests/test_pdb_splitchain.py

@@ -61,6 +61,7 @@ def exec_module(self):
 
         return
 
+
     def test_default(self):
         """$ pdb_splitchain data/dummy.pdb"""
 
@@ -96,6 +97,91 @@ def test_default(self):
 
                 self.assertEqual(fname_chain, list(set(pdb_chains))[0])
 
+    def test_run_fhandler(self):
+        """pdb_splitchain.run(fhandler)"""
+        from pdbtools import pdb_splitchain
+
+        src = os.path.join(data_dir, 'dummy.pdb')
+        dst = os.path.join(self.tempdir, 'dummy.pdb')
+        shutil.copy(src, dst)
+
+        with open(dst, 'r') as fin:
+            pdb_splitchain.run(fin)
+
+        # Read files created by script and then delete
+        ofiles = [f for f in os.listdir(self.tempdir) if f.startswith('dummy')]
+        self.assertEqual(len(ofiles), 4 + 1)  # ori + 4 chains
+
+        # Make sure each file has the chain it should have
+        records = (('ATOM', 'HETATM', 'TER', 'ANISOU'))
+        for fpath in ofiles:
+            if fpath == 'dummy.pdb':
+                continue
+
+            with open(os.path.join(self.tempdir, fpath), 'r') as handle:
+                fname_chain = fpath.split('_')[1][:-4]  # xxx_(X).pdb
+                pdb_chains = [l[21] for l in handle if l.startswith(records)]
+
+                self.assertEqual(fname_chain, list(set(pdb_chains))[0])
+
+    def test_run_iterable(self):
+        """pdb_splitchain.run(iterable)"""
+        from pdbtools import pdb_splitchain
+
+        src = os.path.join(data_dir, 'dummy.pdb')
+        dst = os.path.join(self.tempdir, 'dummy.pdb')
+        shutil.copy(src, dst)
+
+        with open(dst, 'r') as fin:
+            lines = fin.readlines()
+
+        pdb_splitchain.run(lines)
+
+        # Read files created by script and then delete
+        ofiles = [
+            f
+            for f in os.listdir(self.tempdir)
+            if f.startswith('splitchains')
+            ]
+        self.assertEqual(len(ofiles), 4)  # 4 chains
+
+        # Make sure each file has the chain it should have
+        records = (('ATOM', 'HETATM', 'TER', 'ANISOU'))
+        for fpath in ofiles:
+
+            with open(os.path.join(self.tempdir, fpath), 'r') as handle:
+                fname_chain = fpath.split('_')[1][:-4]  # xxx_(X).pdb
+                pdb_chains = [l[21] for l in handle if l.startswith(records)]
+
+                self.assertEqual(fname_chain, list(set(pdb_chains))[0])
+
+    def test_run_iterable_with_name(self):
+        """pdb_splitchain.run(iterable)"""
+        from pdbtools import pdb_splitchain
+
+        src = os.path.join(data_dir, 'dummy.pdb')
+        dst = os.path.join(self.tempdir, 'dummy.pdb')
+        shutil.copy(src, dst)
+
+        with open(dst, 'r') as fin:
+            lines = fin.readlines()
+
+        pdb_splitchain.run(lines, outname='newname')
+
+        # Read files created by script and then delete
+        ofiles = [f for f in os.listdir(self.tempdir) if f.startswith('newname')]
+        self.assertEqual(len(ofiles), 4)  # 4 chains
+
+        # Make sure each file has the chain it should have
+        records = (('ATOM', 'HETATM', 'TER', 'ANISOU'))
+        for fpath in ofiles:
+
+            with open(os.path.join(self.tempdir, fpath), 'r') as handle:
+                fname_chain = fpath.split('_')[1][:-4]  # xxx_(X).pdb
+                pdb_chains = [l[21] for l in handle if l.startswith(records)]
+
+                self.assertEqual(fname_chain, list(set(pdb_chains))[0])
+
     def test_file_not_found(self):
         """$ pdb_splitchain not_existing.pdb"""
 

tests/test_pdb_splitmodel.py

@@ -93,6 +93,88 @@ def test_default(self):
                 n_lines = len(handle.readlines())
                 self.assertEqual(n_lines, 2)
 
+    def test_run_iterable(self):
+        """pdb_splitmodel.run(iterable)"""
+        from pdbtools import pdb_splitmodel
+
+        # Copy input file to tempdir
+
+        # Simulate input
+        src = os.path.join(data_dir, 'ensemble_OK.pdb')
+        dst = os.path.join(self.tempdir, 'ensemble_OK.pdb')
+        shutil.copy(src, dst)
+
+        with open(dst, 'r') as fin:
+            lines = fin.readlines()
+
+        pdb_splitmodel.run(lines)
+
+        # Read files created by script
+        ofiles = [f for f in os.listdir(self.tempdir)
+                  if f.startswith('splitmodels')]
+        self.assertEqual(len(ofiles), 2)
+
+        for fpath in ofiles:
+            if fpath == 'ensemble_OK.pdb':
+                continue
+
+            with open(os.path.join(self.tempdir, fpath), 'r') as handle:
+                n_lines = len(handle.readlines())
+                self.assertEqual(n_lines, 2)
+
+    def test_run_iterable_with_name(self):
+        """pdb_splitmodel.run(iterable)"""
+        from pdbtools import pdb_splitmodel
+
+        # Copy input file to tempdir
+
+        # Simulate input
+        src = os.path.join(data_dir, 'ensemble_OK.pdb')
+        dst = os.path.join(self.tempdir, 'ensemble_OK.pdb')
+        shutil.copy(src, dst)
+
+        with open(dst, 'r') as fin:
+            lines = fin.readlines()
+
+        pdb_splitmodel.run(lines, outname='newname')
+
+        # Read files created by script
+        ofiles = [f for f in os.listdir(self.tempdir)
+                  if f.startswith('newname')]
+        self.assertEqual(len(ofiles), 2)
+
+        for fpath in ofiles:
+            with open(os.path.join(self.tempdir, fpath), 'r') as handle:
+                n_lines = len(handle.readlines())
+                self.assertEqual(n_lines, 2)
+
+    def test_run_fhandler(self):
+        """pdb_splitmodel.run(fhandler)"""
+        from pdbtools import pdb_splitmodel
+
+        # Copy input file to tempdir
+
+        # Simulate input
+        src = os.path.join(data_dir, 'ensemble_OK.pdb')
+        dst = os.path.join(self.tempdir, 'ensemble_OK.pdb')
+        shutil.copy(src, dst)
+
+        with open(dst, 'r') as fin:
+            pdb_splitmodel.run(fin)
+
+        # Read files created by script
+        ofiles = [f for f in os.listdir(self.tempdir)
+                  if f.startswith('ensemble_OK')]
+        self.assertEqual(len(ofiles), 2 + 1)  # ori + 2 models
+
+        for fpath in ofiles:
+            if fpath == 'ensemble_OK.pdb':
+                continue
+
+            with open(os.path.join(self.tempdir, fpath), 'r') as handle:
+                n_lines = len(handle.readlines())
+                self.assertEqual(n_lines, 2)
+
     def test_file_not_found(self):
         """$ pdb_splitmodel not_existing.pdb"""