Made changes based on @rmcgibbo's suggestions. Updated tests to use PDB code…

Made changes based on @rmcgibbo's suggestions.  Updated tests to use PDB code retrieval, switched to simpler test cases for now.

pdbfixer/pdbfixer.py

@@ -58,6 +58,16 @@ else:
 _VALID_URLS = set(uses_relative + uses_netloc + uses_params)
 _VALID_URLS = set(uses_relative + uses_netloc + uses_params)
 _VALID_URLS.discard('')
 _VALID_URLS.discard('')
 
 
+def _is_url(url):
+    """Check to see if a URL has a valid protocol.
+    from pandas/io.common.py Copyright 2014 Pandas Developers
+    Used under the BSD licence
+    """
+    try:
+        return urlparse(url).scheme in _VALID_URLS
+    except:
+        return False
+        
 substitutions = {
 substitutions = {
     '2AS':'ASP', '3AH':'HIS', '5HP':'GLU', 'ACL':'ARG', 'AGM':'ARG', 'AIB':'ALA', 'ALM':'ALA', 'ALO':'THR', 'ALY':'LYS', 'ARM':'ARG',
     '2AS':'ASP', '3AH':'HIS', '5HP':'GLU', 'ACL':'ARG', 'AGM':'ARG', 'AIB':'ALA', 'ALM':'ALA', 'ALO':'THR', 'ALY':'LYS', 'ARM':'ARG',
     'ASA':'ASP', 'ASB':'ASP', 'ASK':'ASP', 'ASL':'ASP', 'ASQ':'ASP', 'AYA':'ALA', 'BCS':'CYS', 'BHD':'ASP', 'BMT':'THR', 'BNN':'ALA',
     'ASA':'ASP', 'ASB':'ASP', 'ASK':'ASP', 'ASL':'ASP', 'ASQ':'ASP', 'AYA':'ALA', 'BCS':'CYS', 'BHD':'ASP', 'BMT':'THR', 'BNN':'ALA',
@@ -157,27 +167,24 @@ class PDBFixer(object):
             pass
             pass
 
 
         if isinstance(structure, str):
         if isinstance(structure, str):
-            # First, try opening it as a local file.
             if os.path.exists(structure):
             if os.path.exists(structure):
+                # First, try opening it as a local file.
                 file = open(structure, 'r')                
                 file = open(structure, 'r')                
                 structure = PdbStructure(file)
                 structure = PdbStructure(file)
                 file.close()
                 file.close()
-            
-            # If it's a URL, try that.
-            if self._is_url(structure):
+            elif _is_url(structure):            
+                # If it's a URL, try that.
                 file = urlopen(structure)
                 file = urlopen(structure)
                 structure = PdbStructure(file)
                 structure = PdbStructure(file)
                 file.close()
                 file.close()
-                
-            # If it's a PDB code, try that.
-            if len(structure) == 4:
+            elif len(structure) == 4:                
+                # If it's a PDB code, try that.
                 url = 'http://www.rcsb.org/pdb/files/%s.pdb' % structure
                 url = 'http://www.rcsb.org/pdb/files/%s.pdb' % structure
                 file = urlopen(url)
                 file = urlopen(url)
                 structure = PdbStructure(file)
                 structure = PdbStructure(file)
                 file.close()
                 file.close()
-                
-        # If a file-like object, read it.
-        if hasattr(structure, 'read'):
+        elif hasattr(structure, 'read'):                
+            # If a file-like object, read it.
             structure = PdbStructure(structure)
             structure = PdbStructure(structure)
 
 
         self.structure = structure
         self.structure = structure
@@ -196,17 +203,6 @@ class PDBFixer(object):
             name = next(templatePdb.topology.residues()).name
             name = next(templatePdb.topology.residues()).name
             self.templates[name] = templatePdb
             self.templates[name] = templatePdb
 
 
-    @classmethod
-    def _is_url(cls, url):
-        """Check to see if a URL has a valid protocol.
-        from pandas/io.common.py Copyright 2014 Pandas Developers
-        Used under the BSD licence
-        """
-        try:
-            return urlparse(url).scheme in _VALID_URLS
-        except:
-            return False
-        
     def _addAtomsToTopology(self, heavyAtomsOnly, omitUnknownMolecules):
     def _addAtomsToTopology(self, heavyAtomsOnly, omitUnknownMolecules):
         """Create a new Topology in which missing atoms have been added."""
         """Create a new Topology in which missing atoms have been added."""
         
         
@@ -740,7 +736,7 @@ def main():
             parser.error('No filename specified')
             parser.error('No filename specified')
         if len(args) > 1:
         if len(args) > 1:
             parser.error('Must specify a single filename')
             parser.error('Must specify a single filename')
-        fixer = PDBFixer(PdbStructure(open(args[0])))
+        fixer = PDBFixer(argv[0])
         if options.residues:
         if options.residues:
             fixer.findMissingResidues()
             fixer.findMissingResidues()
         else:
         else:

tests/test_build_and_simulate.py

 
 
 import pdbfixer
 import pdbfixer
 import simtk.openmm 
 import simtk.openmm 
-import Bio.PDB
 
 
 import os
 import os
 import os.path
 import os.path
@@ -87,29 +86,15 @@ def test_build_and_simulate():
     pdbcodes_to_build = ['1AS5', '1CBN', '1DPO', '1IGY', '1HAG', '1IAO', '4CPA', '1QCQ']
     pdbcodes_to_build = ['1AS5', '1CBN', '1DPO', '1IGY', '1HAG', '1IAO', '4CPA', '1QCQ']
     pdbcodes_to_simulate = ['1AS5', '1CBN', '1DPO', '1IGY', '1HAG', '1IAO', '4CPA', '1QCQ']
     pdbcodes_to_simulate = ['1AS5', '1CBN', '1DPO', '1IGY', '1HAG', '1IAO', '4CPA', '1QCQ']
 
 
-    # Set up PDB retrieval.
-    from Bio.PDB import PDBList
-    pdblist = PDBList(server=PDBList.alternative_download_url)
+    # DEBUG: Simpler test cases.
+    pdbcodes_to_build = ['1AS5', '1VII']
+    pdbcodes_to_simulate = ['1AS5', '1VII']
 
 
     success = True
     success = True
 
 
     for pdbcode in pdbcodes_to_build:
     for pdbcode in pdbcodes_to_build:
         print pdbcode
         print pdbcode
 
 
-        try:
-            # Remove file if it exists already.
-            input_pdb_filename = 'pdb' + pdbcode + '.ent'
-            if os.path.exists(input_pdb_filename):  
-                os.remove(input_pdb_filename)                
-
-            print "Attempting to retrieve PDB code '%s' from %s..." % (pdbcode, PDBList.alternative_download_url)
-            input_pdb_filename = pdblist.retrieve_pdb_file(pdbcode, pdir='.')
-
-        except Exception as e:
-            print str(e)
-            print "Could not download PDB code '%s'" % pdbcode
-            continue
-
         output_pdb_filename = 'output.pdb'
         output_pdb_filename = 'output.pdb'
 
 
         # PDB setup parameters.
         # PDB setup parameters.
@@ -121,18 +106,18 @@ def test_build_and_simulate():
         positiveIon = 'Na+'
         positiveIon = 'Na+'
         negativeIon = 'Cl-'
         negativeIon = 'Cl-'
 
 
-        infile = open(input_pdb_filename)
-        outfile = open(output_pdb_filename, 'w')
         
         
         success = True
         success = True
 
 
+        outfile = open(output_pdb_filename, 'w')
+
         timeout_seconds = 0.1
         timeout_seconds = 0.1
         watchdog = Watchdog(timeout_seconds)
         watchdog = Watchdog(timeout_seconds)
         try:        
         try:        
             from pdbfixer.pdbfixer import PDBFixer, PdbStructure
             from pdbfixer.pdbfixer import PDBFixer, PdbStructure
             from simtk.openmm import app
             from simtk.openmm import app
             stage = "Creating PDBFixer..."
             stage = "Creating PDBFixer..."
-            fixer = PDBFixer(PdbStructure(infile))
+            fixer = PDBFixer(pdbcode)
             stage = "Finding missing residues..."
             stage = "Finding missing residues..."
             fixer.findMissingResidues()
             fixer.findMissingResidues()
             stage = "Finding nonstandard residues..."
             stage = "Finding nonstandard residues..."
@@ -150,7 +135,6 @@ def test_build_and_simulate():
             stage = "Writing PDB file..."
             stage = "Writing PDB file..."
             app.PDBFile.writeFile(fixer.topology, fixer.positions, outfile)
             app.PDBFile.writeFile(fixer.topology, fixer.positions, outfile)
             stage = "Done."
             stage = "Done."
-            infile.close()
             outfile.close()
             outfile.close()
 
 
         except Watchdog:
         except Watchdog:
@@ -182,7 +166,6 @@ def test_build_and_simulate():
             del watchdog
             del watchdog
 
 
         # Clean up.
         # Clean up.
-        os.remove(input_pdb_filename)
         os.remove(output_pdb_filename)
         os.remove(output_pdb_filename)
 
 
     if not success:
     if not success: