Adding WriteTrajMol class for python interface.

pymolfile/pymolfile.py

@@ -309,6 +309,123 @@ class Topology(object):
         #if self.pluginhandle is not None:
         #    libpymolfile.close_file_read(self.pluginhandle)
 
+class WriteMolfile(object):
+
+    def __init__(self, molfile_handle):
+        self.handle = molfile_handle
+
+    def __enter__(self):
+        if self.handle.foutOpen is False:
+            global MOLFILE_PLUGINS
+            global C_MOLFILE_PLUGINS
+            #numlist = libpymolfile.molfile_init()
+            self.handle.wplugin = libpymolfile.get_plugin(C_MOLFILE_PLUGINS, 
+                    self.handle.wfplugin[0])
+            self.handle.wpluginhandle = None
+            if self.handle.topology:
+                self.handle.natoms = self.handle.topology.natoms
+            #if self.handle.trajectory:
+            #    rtn = self.handle.trajectory.iread() 
+            #    if self.handle.trajectory.atoms is not None and rtn is not None:
+            #        self.handle.natoms = len(self.handle.trajectory.atoms["coords"])
+            
+            if self.handle.kwords["silent"]:
+                with stdout_redirected():
+                    self.handle.wpluginhandle = libpymolfile.open_file_write(
+                            self.handle.wplugin, self.handle.wfname, 
+                            self.handle.wtype, self.handle.natoms)
+                    self.handle.foutOpen = True
+            else:
+                self.handle.wpluginhandle = libpymolfile.open_file_write(
+                        self.handle.wplugin, self.handle.wfname, 
+                        self.handle.wtype, self.handle.natoms)
+                self.handle.foutOpen = True
+
+        if self.handle.wpluginhandle is not None:
+            return self
+        else:
+            return None
+
+    def __exit__(self, type, value, traceback):
+        if self.handle.foutEOF:
+            if self.handle.foutOpen:
+                if self.handle.wpluginhandle is not None:
+                    rtn = libpymolfile.close_file_write(self.handle.wpluginhandle)
+                    if rtn:
+                        self.handle.wpluginhandle = None
+    
+    def write_topology(self):
+        status = False
+        bstatus = False
+        astatus = False
+        if self.handle.topology.bonds is not None:
+            if self.handle.wfplugin[1][7]>0:
+                bstatus = libpymolfile.write_fill_bonds(self.handle.wpluginhandle, self.handle.topology.bonds)
+        if self.handle.topology.angles is not None:
+            if self.handle.wfplugin[1][8]>0:
+                astatus = libpymolfile.write_fill_angles(self.handle.wpluginhandle, self.handle.topology.angles)
+        if self.handle.wfplugin[1][6]>0 and self.handle.topology.structure is not None:
+            prototype = molfile_prototype(self.handle.wtype)
+            newstructure = [
+                    (str(item[0]),
+                     str(item[1]),
+                     str(item[2]),
+                     int(item[3]),
+                     str(item[4]),
+                     str(item[5]),
+                     str(item[6]),
+                     str(item[7]),
+                     float(item[8]),
+                     float(item[9]),
+                     float(item[10]),
+                     float(item[11]),
+                     float(item[12]),
+                     int(item[13])
+                     ) if len(item)<15 else ( 
+                         str(item[0]),
+                         str(item[1]),
+                         str(item[2]),
+                         int(item[3]),
+                         str(item[4]),
+                         str(item[5]),
+                         str(item[6]),
+                         str(item[7]),
+                         float(item[8]),
+                         float(item[9]),
+                         float(item[10]),
+                         float(item[11]),
+                         float(item[12]),
+                         int(item[13]),
+                         int(item[14])
+                         ) for item in self.handle.topology.structure
+                     ]
+            thestructure = np.array(newstructure, dtype=prototype.dtype)
+            status = libpymolfile.write_fill_structure(self.handle.wpluginhandle, thestructure)
+        else:
+            status = True
+        return status
+
+    def write_trajectory(self):
+        if self.handle.wfplugin[1][9]>0 and self.handle.trajectory:
+            rtnOut = None
+            done = False
+            print("Writing Trajectory")
+            if self.handle.trajectory.atoms is not None:
+                rtnOut = libpymolfile.write_fill_timestep(self.handle.wpluginhandle, self.handle.trajectory.atoms)
+            while done is False:
+                rtn = self.handle.trajectory.iread() 
+                if self.handle.trajectory.atoms is not None and rtn is not None:
+                    print("Writing Trajectory")
+                    rtnOut = libpymolfile.write_fill_timestep(self.handle.wpluginhandle, self.handle.trajectory.atoms)
+                else:
+                    print("Done")
+                    done = True
+                    self.handle.foutEOF = True
+            return True
+        else:
+            self.handle.foutEOF = True
+            return True
+
 class OpenMolTraj(object):
 
     def __init__(self, traj_handle):
@@ -496,8 +613,12 @@ def read_topology(file_name, file_format, plugin, silent):
     if(topo.structure is not None or 
        topo.bonds is not None or
        topo.angles is not None):
+        if topo.pluginhandle is not None:
+            libpymolfile.close_file_read(topo.pluginhandle)
         return topo
     else:
+        if topo.pluginhandle is not None:
+            libpymolfile.close_file_read(topo.pluginhandle)
         del topo
         return None
 
@@ -651,11 +772,39 @@ class OpenMolfile(object):
                     warnings.warn("Pymolfile can not find a plugin to open the '" + self.kwords["file_format"] + 
                             "' file format of the file " + self.kwords["file_name"]) 
 
+    def save(self, wfname, wtype=None):
+        self.wfname = wfname
+        if wtype is None:
+            file_dir, file_base, file_ext = get_dir_base_extension(wfname)
+            if file_ext:
+                self.wtype = file_ext
+            else:
+                self.wtype = file_base
+        else:
+            self.wtype = wtype
+        self.wfplugin = get_plugin_with_ext(wtype)
+        if self.wfplugin:
+            done = False
+            structOK = False
+            with WriteMolfile(self) as f:
+                if f:
+                    if self.wfplugin[1][6]>0:
+                        structOK = f.write_topology()
+                    if self.wfplugin[1][9]>0:
+                        done = f.write_trajectory()
+
     def initialize_settings(self):
+        self.foutEOF = False
+        self.foutOpen = False
         self.trajectory = None
         self.topology = None
         self.fplugin = None
         self.tplugin = None
+        self.wfname = None
+        self.wtype = None
+        self.wplugin = None
+        self.wfplugin = None
+        self.natoms = None
         self.smolplugin = None
         self.cmolplugin = None
         self.kwords = {