PyLammps: alternative OutputCapture using tempfiles

python/lammps/pylammps.py

@@ -20,47 +20,47 @@
 
 from __future__ import print_function
 
+import io
 import os
 import re
-import select
+import sys
+import tempfile
 from collections import namedtuple
 
 from .core import lammps
 
+# -------------------------------------------------------------------------
 
 class OutputCapture(object):
   """ Utility class to capture LAMMPS library output """
-
   def __init__(self):
-    self.stdout_pipe_read, self.stdout_pipe_write = os.pipe()
+    self.stdout_fd = sys.stdout.fileno()
-    self.stdout_fd = 1
+    self.captured_output = ""
 
   def __enter__(self):
-    self.stdout = os.dup(self.stdout_fd)
+    self.tmpfile = tempfile.TemporaryFile(mode='w+b')
-    os.dup2(self.stdout_pipe_write, self.stdout_fd)
+
+    sys.stdout.flush()
+
+    # make copy of original stdout
+    self.stdout_orig = os.dup(self.stdout_fd)
+
+    # replace stdout and redirect to temp file
+    os.dup2(self.tmpfile.fileno(), self.stdout_fd)
     return self
 
   def __exit__(self, exc_type, exc_value, traceback):
-    os.dup2(self.stdout, self.stdout_fd)
+    os.dup2(self.stdout_orig, self.stdout_fd)
-    os.close(self.stdout)
+    os.close(self.stdout_orig)
-    os.close(self.stdout_pipe_read)
+    self.tmpfile.close()
-    os.close(self.stdout_pipe_write)
-
-  # check if we have more to read from the pipe
-  def more_data(self, pipe):
-    r, _, _ = select.select([pipe], [], [], 0)
-    return bool(r)
-
-  # read the whole pipe
-  def read_pipe(self, pipe):
-    out = ""
-    while self.more_data(pipe):
-      out += os.read(pipe, 1024).decode()
-    return out
 
   @property
   def output(self):
-    return self.read_pipe(self.stdout_pipe_read)
+    sys.stdout.flush()
+    self.tmpfile.flush()
+    self.tmpfile.seek(0, io.SEEK_SET)
+    self.captured_output = self.tmpfile.read().decode('utf-8')
+    return self.captured_output
 
 # -------------------------------------------------------------------------
 
@@ -109,9 +109,9 @@ class AtomList(object):
     """
     if index not in self._loaded:
         if self.dimensions == 2:
-            atom = Atom2D(self._pylmp, index + 1)
+            atom = Atom2D(self._pylmp, index)
         else:
-            atom = Atom(self._pylmp, index + 1)
+            atom = Atom(self._pylmp, index)
         self._loaded[index] = atom
     return self._loaded[index]
 
@@ -612,7 +612,8 @@ class PyLammps(object):
     :getter: Returns an object with properties storing the current system state
     :type: namedtuple
     """
-    output = self.info("system")
+    output = self.lmp_info("system")
+    output = output[output.index("System information:")+1:]
     d = self._parse_info_system(output)
     return namedtuple('System', d.keys())(*d.values())
 
@@ -624,7 +625,8 @@ class PyLammps(object):
     :getter: Returns an object with properties storing the current communication state
     :type: namedtuple
     """
-    output = self.info("communication")
+    output = self.lmp_info("communication")
+    output = output[output.index("Communication information:")+1:]
     d = self._parse_info_communication(output)
     return namedtuple('Communication', d.keys())(*d.values())
 
@@ -636,7 +638,8 @@ class PyLammps(object):
     :getter: Returns a list of computes that are currently active in this LAMMPS instance
     :type: list
     """
-    output = self.info("computes")
+    output = self.lmp_info("computes")
+    output = output[output.index("Compute information:")+1:]
     return self._parse_element_list(output)
 
   @property
@@ -647,7 +650,8 @@ class PyLammps(object):
     :getter: Returns a list of dumps that are currently active in this LAMMPS instance
     :type: list
     """
-    output = self.info("dumps")
+    output = self.lmp_info("dumps")
+    output = output[output.index("Dump information:")+1:]
     return self._parse_element_list(output)
 
   @property
@@ -658,7 +662,8 @@ class PyLammps(object):
     :getter: Returns a list of fixes that are currently active in this LAMMPS instance
     :type: list
     """
-    output = self.info("fixes")
+    output = self.lmp_info("fixes")
+    output = output[output.index("Fix information:")+1:]
     return self._parse_element_list(output)
 
   @property
@@ -669,7 +674,8 @@ class PyLammps(object):
     :getter: Returns a list of atom groups that are currently active in this LAMMPS instance
     :type: list
     """
-    output = self.info("groups")
+    output = self.lmp_info("groups")
+    output = output[output.index("Group information:")+1:]
     return self._parse_groups(output)
 
   @property
@@ -680,11 +686,12 @@ class PyLammps(object):
     :getter: Returns a dictionary of all variables that are defined in this LAMMPS instance
     :type: dict
     """
-    output = self.info("variables")
+    output = self.lmp_info("variables")
-    vars = {}
+    output = output[output.index("Variable information:")+1:]
+    variables = {}
     for v in self._parse_element_list(output):
-      vars[v['name']] = Variable(self, v['name'], v['style'], v['def'])
+      variables[v['name']] = Variable(self, v['name'], v['style'], v['def'])
-    return vars
+    return variables
 
   def eval(self, expr):
     """
@@ -709,10 +716,9 @@ class PyLammps(object):
     return [x.strip() for x in value.split('=')]
 
   def _parse_info_system(self, output):
-    lines = output[5:-2]
     system = {}
 
-    for line in lines:
+    for line in output:
       if line.startswith("Units"):
         system['units'] = self._get_pair(line)[1]
       elif line.startswith("Atom style"):
@@ -770,10 +776,9 @@ class PyLammps(object):
     return system
 
   def _parse_info_communication(self, output):
-    lines = output[5:-3]
     comm = {}
 
-    for line in lines:
+    for line in output:
       if line.startswith("MPI library"):
         comm['mpi_version'] = line.split(':')[1].strip()
       elif line.startswith("Comm style"):
@@ -791,10 +796,10 @@ class PyLammps(object):
     return comm
 
   def _parse_element_list(self, output):
-    lines = output[5:-3]
     elements = []
 
-    for line in lines:
+    for line in output:
+      if not line or (":" not in line): continue
       element_info = self._split_values(line.split(':')[1].strip())
       element = {'name': element_info[0]}
       for key, value in [self._get_pair(x) for x in element_info[1:]]:
@@ -803,11 +808,10 @@ class PyLammps(object):
     return elements
 
   def _parse_groups(self, output):
-    lines = output[5:-3]
     groups = []
     group_pattern = re.compile(r"(?P<name>.+) \((?P<type>.+)\)")
 
-    for line in lines:
+    for line in output:
       m = group_pattern.match(line.split(':')[1].strip())
       group = {'name': m.group('name'), 'type': m.group('type')}
       groups.append(group)
@@ -830,6 +834,15 @@ class PyLammps(object):
     'thermo_modify', 'thermo_style', 'timestep', 'undump', 'unfix', 'units',
     'variable', 'velocity', 'write_restart'] + self.lmp.available_styles("command")))
 
+  def lmp_info(self, s):
+      # skip anything before and after Info-Info-Info
+      # also skip timestamp line
+      output = self.__getattr__("info")(s)
+      indices = [index for index, line in enumerate(output) if line.startswith("Info-Info-Info-Info")]
+      start = indices[0]
+      end = indices[1]
+      return [line for line in output[start+2:end] if line]
+
   def __getattr__(self, name):
     """
     This method is where the Python 'magic' happens. If a method is not