Add PyLammps.atoms test and fix bug with -echo screen/both

python/lammps.py

@@ -1974,6 +1974,14 @@ class PyLammps(object):
   """
 
   def __init__(self, name="", cmdargs=None, ptr=None, comm=None):
+    self.has_echo = False
+
+    if cmdargs:
+      if '-echo' in cmdargs:
+        idx = cmdargs.index('-echo')
+        # ensures that echo line is ignored during output capture
+        self.has_echo = idx+1 < len(cmdargs) and cmdargs[idx+1] in ('screen', 'both')
+
     if ptr:
       if isinstance(ptr,PyLammps):
         self.lmp = ptr.lmp
@@ -2323,7 +2331,8 @@ class PyLammps(object):
       cmd_args = [name] + [str(x) for x in args]
 
       with OutputCapture() as capture:
-        self.command(' '.join(cmd_args))
+        cmd = ' '.join(cmd_args)
+        self.command(cmd)
         output = capture.output
 
       if 'verbose' in kwargs and kwargs['verbose']:
@@ -2331,6 +2340,9 @@ class PyLammps(object):
 
       lines = output.splitlines()
 
+      if self.has_echo:
+        lines = lines[1:]
+
       if len(lines) > 1:
         return lines
       elif len(lines) == 1:

unittest/python/python-pylammps.py

@@ -7,6 +7,9 @@ class PythonPyLammps(unittest.TestCase):
         if 'LAMMPS_MACHINE_NAME' in os.environ:
             machine=os.environ['LAMMPS_MACHINE_NAME']
         self.pylmp = PyLammps(name=machine,  cmdargs=['-nocite', '-log','none', '-echo', 'screen'])
+        self.pylmp.units("lj")
+        self.pylmp.atom_style("atomic")
+        self.pylmp.atom_modify("map array")
 
         if 'LAMMPS_CMAKE_CACHE' in os.environ:
             self.cmake_cache = {}
@@ -45,6 +48,9 @@ class PythonPyLammps(unittest.TestCase):
 
         self.assertEqual(self.pylmp.lmp.create_atoms(2, id=None, type=types, x=x), 2)
         self.assertEqual(self.pylmp.system.natoms, 2)
+        self.assertEqual(len(self.pylmp.atoms), 2)
+        self.assertEqual(self.pylmp.atoms[0].position, tuple(x[0:3]))
+        self.assertEqual(self.pylmp.atoms[1].position, tuple(x[3:6]))
 
 if __name__ == "__main__":
     unittest.main()