python: update examples and docs

python/lammps/ipython/magics.py

@@ -0,0 +1,75 @@
+# ----------------------------------------------------------------------
+#   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
+#   https://www.lammps.org/ Sandia National Laboratories
+#   LAMMPS Development team: developers@lammps.org
+#
+#   Copyright (2003) Sandia Corporation.  Under the terms of Contract
+#   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
+#   certain rights in this software.  This software is distributed under
+#   the GNU General Public License.
+#
+#   See the README file in the top-level LAMMPS directory.
+# -------------------------------------------------------------------------
+
+################################################################################
+# IPython/Jupyter Notebook additions
+# Written by Richard Berger <richard.berger@outlook.com>
+################################################################################
+
+import io
+import os
+import sys
+import tempfile
+from IPython.core.magic import (Magics, magics_class, cell_magic)
+import IPython.core.magic_arguments as magic_arguments
+
+class OutputCapture(object):
+  """ Utility class to capture LAMMPS library output """
+  def __init__(self):
+    self.stdout_fd = 1
+    self.captured_output = ""
+
+  def __enter__(self):
+    self.tmpfile = tempfile.TemporaryFile(mode='w+b')
+
+    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_orig, self.stdout_fd)
+    os.close(self.stdout_orig)
+    self.tmpfile.close()
+
+  @property
+  def output(self):
+    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
+
+# -------------------------------------------------------------------------
+
+@magics_class
+class LammpsMagics(Magics):
+    @magic_arguments.magic_arguments()
+    @magic_arguments.argument('output', type=str, default='', nargs='?',
+        help="""The name of the variable in which to store output.
+
+        If unspecified, captured output is discarded.
+        """
+    )
+    @cell_magic
+    def capture_lammps_output(self, line, cell):
+        """run the cell, capturing LAMMPS stdout and stderr."""
+        args = magic_arguments.parse_argstring(self.capture_lammps_output, line)
+        with OutputCapture() as capture:
+            self.shell.run_cell(cell)
+            if args.output:
+                self.shell.user_ns[args.output] = capture.output