Merge pull request #69 from drorlab/68-recognize-single-string-trj-ex…

…tract_coordinates

make extract_coordinates work with a single string as trajectory input

pensa/preprocessing/coordinates.py

@@ -18,8 +18,8 @@ def extract_coordinates(top, pdb, trj_list, out_name, sel_string, start_frame=0,
             Can read all MDAnalysis-compatible topology formats.
         pdb : str
             File name for the reference PDB file.
-        trj_list : list of str
-            File names for the input trajectory.
+        trj_list : str or list of str
+            File name(s) for the input trajectory.
             Can read all MDAnalysis-compatible trajectory formats.
         out_name : str
             Core of the file names for the output files.
@@ -29,6 +29,9 @@ def extract_coordinates(top, pdb, trj_list, out_name, sel_string, start_frame=0,
             First frame to read from the trajectory.
 
     """
+    # Convert the trajectory input to a list if it is a single string
+    if type(trj_list) is str:
+        trj_list = [trj_list]
     # Read the topology+PDB files and extract selected parts.
     u = mda.Universe(top, pdb)
     u.residues.resids -= residues_offset

tests/test_workflow.py

@@ -100,14 +100,17 @@ def setUp(self):
         )
 
         # Extract the coordinates of the TM region from the trajectory
+        # - Test using a single trajectory directly
         self.file_tm_single_a = extract_coordinates(
-            self.ref_file_a, self.pdb_file_a, [self.trj_file_a],
+            self.ref_file_a, self.pdb_file_a, self.trj_file_a,
             self.trj_name_a + "_tm", sel_string_a
         )
+        # - Test using a single trajectory as a list
         self.file_tm_single_b = extract_coordinates(
             self.ref_file_b, self.pdb_file_b, [self.trj_file_b],
             self.trj_name_b + "_tm", sel_string_b
         )
+        # - Test using multiple trajectories
         self.file_tm_combined = extract_coordinates_combined(
             [self.ref_file_a] * 3 + [self.ref_file_b] * 3,
             self.trj_file_a + self.trj_file_b,