Adding functions for structure and trajectory writing through molfile.save

pymolfile/molfile/libpymolfile.i

@@ -20,7 +20,7 @@
 
 %{
 /* Python SWIG interface to libpymolfile
-   Copyright (c) 2017 Berk Onat <b.onat@warwick.ac.uk>
+   Copyright (c) 2018 Berk Onat <b.onat@warwick.ac.uk>
    Published with UIUC LICENSE
  */
 #define SWIG_FILE_WITH_INIT
@@ -79,7 +79,7 @@ extern PyObject * molfile_plugin_info(PyObject* molcapsule, int plugin_no);
 %inline %{
 PyObject * my_open_file_read(PyObject* molcapsule, char* fname, char* ftype, int natoms) {
     if (PyType_Ready(&MolObjectType) < 0)
-        return NULL;
+        Py_RETURN_NONE;
     PyTypeObject *type = &MolObjectType;
     MolObject *plugin_c;
     molfile_plugin_t* plugin = (molfile_plugin_t*) PyMolfileCapsule_AsVoidPtr(molcapsule);
@@ -97,6 +97,82 @@ PyObject * my_open_file_read(PyObject* molcapsule, char* fname, char* ftype, int
   }
 %}
 
+%feature("autodoc", "0") my_open_file_write;
+%rename (open_file_write) my_open_file_write;
+%exception my_open_file_write {
+  $action
+  if (PyErr_Occurred()) SWIG_fail;
+}
+%inline %{
+PyObject * my_open_file_write(PyObject* molcapsule, char* fname, char* ftype, int natoms) {
+    if (PyType_Ready(&MolObjectType) < 0)
+        Py_RETURN_NONE;
+    if(natoms < 1)
+        Py_RETURN_NONE;
+    PyTypeObject *type = &MolObjectType;
+    MolObject *plugin_c;
+    molfile_plugin_t* plugin = (molfile_plugin_t*) PyMolfileCapsule_AsVoidPtr(molcapsule);
+    plugin_c = (MolObject *)type->tp_alloc(type, 0);
+    plugin_c->plugin = molcapsule;
+    void *file_handle = plugin->open_file_write(fname, ftype, natoms);
+    plugin_c->file_handle = PyMolfileCapsule_FromVoidPtr(file_handle, del_molfile_file_handle);
+    if (!plugin_c->file_handle) {
+        Py_RETURN_NONE;
+    } else {
+        plugin_c->natoms = natoms;
+        return (PyObject *)plugin_c;
+    }
+  }
+%}
+
+/*
+%feature("autodoc", "0") my_set_file_write_stdout;
+%rename (set_file_write_stdout) my_set_file_write_stdout;
+%exception my_set_file_write_stdout {
+  $action
+  if (PyErr_Occurred()) SWIG_fail;
+}
+%inline %{
+PyObject * my_set_file_write_stdout(PyObject* molpack) {
+    MolObject* plugin_handle = (MolObject*) molpack;
+    PyObject* filecapsule = plugin_handle->file_handle;
+    void *file_handle = (void*) PyMolfileCapsule_AsVoidPtr(filecapsule);
+    if(file_handle->fd){
+        fclose(file_handle->fd);
+        file_handle->fd = stdout;
+    }
+    else if(file_handle->fp){
+        fclose(file_handle->fp);
+        file_handle->fp = stdout;
+    }
+    else if(file_handle->file){
+        fclose(file_handle->file);
+        file_handle->file = stdout;
+    }
+    else if(file_handle->mf){
+        fclose(file_handle->mf);
+        file_handle->mf = stdout;
+    }
+    else if(file_handle->writer){
+        if(file_handle->writer->fd){
+            fclose(file_handle->writer->fd);
+            file_handle->writer->fd = stdout;
+        }
+    }
+    else (file_handle){
+        fclose(file_handle);
+        file_handle = stdout;
+    }
+    plugin_handle->file_handle = PyMolfileCapsule_FromVoidPtr(file_handle, del_molfile_file_handle);
+    if (!plugin_c->file_handle) {
+        Py_RETURN_NONE;
+    } else {
+        return (PyObject *)plugin_handle;
+    }
+  }
+%}
+*/
+
 %feature("autodoc", "0") my_close_file_read;
 %rename (close_file_read) my_close_file_read;
 %exception my_close_file_read {
@@ -115,18 +191,49 @@ PyObject * my_close_file_read(PyObject* molpack) {
   }
 %}
 
+%feature("autodoc", "0") my_close_file_write;
+%rename (close_file_write) my_close_file_write;
+%exception my_close_file_write {
+  $action
+  if (PyErr_Occurred()) SWIG_fail;
+}
+%inline %{
+PyObject * my_close_file_write(PyObject* molpack) {
+    MolObject* plugin_handle = (MolObject*) molpack;
+    PyObject* plugincapsule = plugin_handle->plugin;   
+    PyObject* filecapsule = plugin_handle->file_handle;
+    molfile_plugin_t* plugin = (molfile_plugin_t*) PyMolfileCapsule_AsVoidPtr(plugincapsule);
+    void *file_handle = (void*) PyMolfileCapsule_AsVoidPtr(filecapsule);
+    plugin->close_file_write(file_handle);
+    Py_RETURN_TRUE;
+  }
+%}
+
+
 %feature("autodoc", "0") read_fill_structure;
 extern PyObject * read_fill_structure(PyObject* molpack, PyObject* prototype);
 
+%feature("autodoc", "0") write_fill_structure;
+extern PyObject * write_fill_structure(PyObject* molpack, PyObject* molarray);
+
 %feature("autodoc", "0") read_fill_bonds;
 extern PyObject * read_fill_bonds(PyObject* molpack);
 
+%feature("autodoc", "0") write_fill_bonds;
+extern PyObject * write_fill_bonds(PyObject* molpack, PyObject* moldict);
+
 %feature("autodoc", "0") read_fill_angles;
 extern PyObject * read_fill_angles(PyObject* molpack);
 
+%feature("autodoc", "0") write_fill_angles;
+extern PyObject * write_fill_angles(PyObject* molpack, PyObject* moldict);
+
 %feature("autodoc", "0") read_fill_next_timestep;
 extern PyObject * read_fill_next_timestep(PyObject* molpack);
 
+%feature("autodoc", "0") write_fill_timestep;
+extern PyObject * write_fill_timestep(PyObject* molpack, PyObject* moldict);
+
 %feature("autodoc", "0") are_plugins_same;
 PyObject* are_plugins_same(PyObject* molpack_a, PyObject* molpack_b);
 

pymolfile/molfile/pymolfile.h

@@ -263,9 +263,13 @@ static PyTypeObject MolObjectType = {
 
 PyObject * molfile_plugin_info(PyObject* molcapsule, int plugin_no);
 PyObject * read_fill_structure(PyObject* molpack, PyObject* prototype);
+PyObject * write_fill_structure(PyObject* molpack, PyObject* molarray);
 PyObject * read_fill_bonds(PyObject* molpack);
+PyObject * write_fill_bonds(PyObject* molpack, PyObject* moldict);
 PyObject * read_fill_angles(PyObject* molpack);
+PyObject * write_fill_angles(PyObject* molpack, PyObject* moldict);
 PyObject * read_fill_next_timestep(PyObject* molpack);
+PyObject * write_fill_timestep(PyObject* molpack, PyObject* moldict);
 PyObject * are_plugins_same(PyObject* molpack_a, PyObject* molpack_b);
 PyObject * are_filehandles_same(PyObject* molpack_a, PyObject* molpack_b);
 PyObject * get_plugin(PyObject* molcapsule, int plug_no);

pymolfile/pymolfile.py

@@ -7,6 +7,9 @@ import re
 import sys
 import numpy as np
 from contextlib import contextmanager
+import platform
+import ctypes
+import tempfile
 import warnings
 
 if sys.version_info > (3,):
@@ -19,6 +22,18 @@ except ImportError:
 
 from .plugin_list import plugins, byte_str_decode, MOLFILE_PLUGINS, C_MOLFILE_PLUGINS
 
+if platform.system() == "Windows":
+    libc = ctypes.cdll.msvcrt
+else:
+    libc = ctypes.CDLL(None)
+
+if platform.system() == "Windows":
+    c_stdout = ctypes.c_void_p.in_dll(libc, 'stdout')
+elif "Darwin" in platform.system():
+    c_stdout = ctypes.c_void_p.in_dll(libc, '__stdoutp')
+else:
+    c_stdout = ctypes.c_void_p.in_dll(libc, 'stdout')
+
 def list_plugins():
     global MOLFILE_PLUGINS
     return MOLFILE_PLUGINS[:]
@@ -49,6 +64,8 @@ def stdout_redirected(to=os.devnull):
     ####assert libc.fileno(ctypes.c_void_p.in_dll(libc, "stdout")) == fd == 1
 
     def _redirect_stdout(to):
+        # Flush the C-level buffer stdout
+        libc.fflush(c_stdout)
         sys.stdout.close() # + implicit flush()
         os.dup2(to.fileno(), fd) # fd writes to 'to' file
         sys.stdout = os.fdopen(fd, 'w') # Python writes to fd
@@ -63,6 +80,42 @@ def stdout_redirected(to=os.devnull):
                                             # buffering and flags such as
                                             # CLOEXEC may be different
 
+# The function stdout_redirect_stream is taken from 
+#https://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/
+
+@contextmanager
+def stdout_redirect_stream(stream):
+    # The original fd stdout points to. Usually 1 on POSIX systems.
+    original_stdout_fd = sys.stdout.fileno()
+
+    def _redirect_stdout(to_fd):
+        """Redirect stdout to the given file descriptor."""
+        # Flush the C-level buffer stdout
+        libc.fflush(c_stdout)
+        # Flush and close sys.stdout - also closes the file descriptor (fd)
+        sys.stdout.close()
+        # Make original_stdout_fd point to the same file as to_fd
+        os.dup2(to_fd, original_stdout_fd)
+        # Create a new sys.stdout that points to the redirected fd
+        sys.stdout = io.TextIOWrapper(os.fdopen(original_stdout_fd, 'wb'))
+
+    # Save a copy of the original stdout fd in saved_stdout_fd
+    saved_stdout_fd = os.dup(original_stdout_fd)
+    try:
+        # Create a temporary file and redirect stdout to it
+        tfile = tempfile.TemporaryFile(mode='w+b')
+        _redirect_stdout(tfile.fileno())
+        # Yield to caller, then redirect stdout back to the saved fd
+        yield
+        _redirect_stdout(saved_stdout_fd)
+        # Copy contents of temporary file to the given stream
+        tfile.flush()
+        tfile.seek(0, io.SEEK_SET)
+        stream.write(tfile.read())
+    finally:
+        tfile.close()
+        os.close(saved_stdout_fd)
+
 def get_dir_base_extension(file_name):
     """ Splits directory, file base and file extensions
 
@@ -346,6 +399,64 @@ class Trajectory(object):
         finally:
             del iter_obj
 
+def molfile_prototype(file_format=None):
+    """ Generates a prototype numpy array that has the same struct 
+        definition as in molfile_atom_t arrays.
+
+        Returns: numpy.array
+    """
+    if file_format is None:
+        file_format = ''
+    #if 0 && vmdplugin_ABIVERSION > 17
+    #  /* The new PDB file formats allows for much larger structures, */
+    #  /* which can therefore require longer chain ID strings.  The   */
+    #  /* new PDBx/mmCIF file formats do not have length limits on    */
+    #  /* fields, so PDB chains could be arbitrarily long strings     */
+    #  /* in such files.  At present, we know we need at least 3-char */
+    #  /* chains for existing PDBx/mmCIF files.                       */
+    #  char chain[4];      /**< required chain name, or ""            */
+    #else
+    #  char chain[2];      /**< required chain name, or ""            */
+    #endif
+    #
+    # Change 'chain', S2 to S4
+    #
+    #if('pdb' in file_format or
+    #   'psf' in file_format):
+    chain_size = 'S4'
+    #else:
+    #    chain_size = 'S2'
+    if('dtr' in file_format.lower() or 
+       'stk' in file_format.lower() or 
+       'atr' in file_format.lower()):
+       prototype = np.array([
+           ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6,0), 
+           ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6,0)
+           ],
+           dtype=[
+               ('name', 'S16'), ('type', 'S16'), ('resname', 'S8'),
+               ('resid', 'i4'), ('segid', 'S8'), ('chain', chain_size),
+               ('altloc', 'S2'), ('insertion', 'S2'), ('occupancy', 'f4'),
+               ('bfactor', 'f4'), ('mass', 'f4'), ('charge', 'f4'),
+               ('radius', 'f4'), ('atomicnumber', 'i4'), ('ctnumber', 'i4')
+               ]
+           )
+    else:
+       prototype = np.array([
+           ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6),
+           ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6)
+           ],
+           dtype=[
+               ('name', 'S16'), ('type', 'S16'), ('resname', 'S8'),
+               ('resid', 'i4'), ('segid', 'S8'), ('chain', chain_size),
+               ('altloc', 'S2'), ('insertion', 'S2'), ('occupancy', 'f4'),
+               ('bfactor', 'f4'), ('mass', 'f4'), ('charge', 'f4'),
+               ('radius', 'f4'), ('atomicnumber', 'i4')
+               ]
+           )
+
+    return prototype
+
 def read_topology(file_name, file_format, plugin, silent):
     """ Reads structure, bonds, angles, dihedrals, impropers and 
         additional informations through molfile_plugin if the 
@@ -373,53 +484,7 @@ def read_topology(file_name, file_format, plugin, silent):
        plugin is not None):
         natoms=0
         topo = Topology(plugin, file_name, file_format, natoms, silent)
-        #if 0 && vmdplugin_ABIVERSION > 17
-        #  /* The new PDB file formats allows for much larger structures, */
-        #  /* which can therefore require longer chain ID strings.  The   */
-        #  /* new PDBx/mmCIF file formats do not have length limits on    */
-        #  /* fields, so PDB chains could be arbitrarily long strings     */
-        #  /* in such files.  At present, we know we need at least 3-char */
-        #  /* chains for existing PDBx/mmCIF files.                       */
-        #  char chain[4];      /**< required chain name, or ""            */
-        #else
-        #  char chain[2];      /**< required chain name, or ""            */
-        #endif
-        #
-        # Change 'chain', S2 to S4
-        #
-        #if('pdb' in file_format or
-        #   'psf' in file_format):
-        chain_size = 'S4'
-        #else:
-        #    chain_size = 'S2'
-        if('dtr' in file_format or 
-           'stk' in file_format or 
-           'atr' in file_format):
-            prototype = np.array([
-                ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6,0), 
-                ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6,0)
-                ],
-                dtype=[
-                    ('name', 'S16'), ('type', 'S16'), ('resname', 'S8'),
-                    ('resid', 'i4'), ('segid', 'S8'), ('chain', chain_size),
-                    ('altloc', 'S2'), ('insertion', 'S2'), ('occupancy', 'f4'),
-                    ('bfactor', 'f4'), ('mass', 'f4'), ('charge', 'f4'),
-                    ('radius', 'f4'), ('atomicnumber', 'i4'), ('ctnumber', 'i4')
-                    ]
-                )
-        else:
-            prototype = np.array([
-                ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6),
-                ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6)
-                ],
-                dtype=[
-                    ('name', 'S16'), ('type', 'S16'), ('resname', 'S8'),
-                    ('resid', 'i4'), ('segid', 'S8'), ('chain', chain_size),
-                    ('altloc', 'S2'), ('insertion', 'S2'), ('occupancy', 'f4'),
-                    ('bfactor', 'f4'), ('mass', 'f4'), ('charge', 'f4'),
-                    ('radius', 'f4'), ('atomicnumber', 'i4')
-                    ]
-                )
+        prototype = molfile_prototype(file_format)
 
         if topo.read_structure(prototype) is not None:
             topo.structure = decode_array(topo._structure)

setup.py

@@ -98,7 +98,7 @@ def check_tcl_version():
         val = None
     return val
 
-VERSION = "0.0.4"
+VERSION = "0.0.5"
 CLASSIFIERS = [
     "Development Status :: 1 - Alpha",
     "Intended Audience :: Science/Research",