##############################################################################
# Institute for the Design of Advanced Energy Systems Process Systems
# Engineering Framework (IDAES PSE Framework) Copyright (c) 2018, by the
# software owners: The Regents of the University of California, through
# Lawrence Berkeley National Laboratory, National Technology & Engineering
# Solutions of Sandia, LLC, Carnegie Mellon University, West Virginia
# University Research Corporation, et al. All rights reserved.
#
# Please see the files COPYRIGHT.txt and LICENSE.txt for full copyright and
# license information, respectively. Both files are also available online
# at the URL "https://github.com/IDAES/idaes".
##############################################################################
"""
Utility functions, etc. for testing
"""
# stdlib
import glob
import os
import re
import signal
import subprocess
import time
import warnings
# third-party
import nbformat
from nbconvert.preprocessors import ExecutePreprocessor
import psutil
# idaes
import idaes
[docs]class RunnableExample(object):
"""Encapsulate a runnable example script.
Example usage (e.g. in a test):
ex = RunnableExample("dmf/resource_example.py")
rcode = ex.run(timeout=5)
assert rcode == 0, "dmf/resource_example.py failed"
"""
ansi_escape = re.compile(r'\x1B\[[0-?]*[ -/]*[@-~]')
def __init__(self, path):
"""Create new example with given script path.
Args:
path(str): Relative (to 'examples/' dir) or absolute
path to the script.
"""
self._path = path
self._reset()
def _reset(self):
self._status = -1
self._out, self._err = '', ''
[docs] def run(self, args=None, timeout=15):
"""Run the example.
* Files ending in '.py' are run with the Python interpreter.
* Files ending in '.ipynb' (Jupyter notebooks) are executed using
Jupyter nbconvert's pre-processing API.
* All other files are simply executed.
Args:
args (list[str]): Arguments passed to scripts (except notebooks)
timeout (int): Time, in seconds, to wait for it to run.
Returns:
Numeric status code
"""
self._reset()
args = [] if args is None else args
script = self._get_script_path()
if script.endswith('.ipynb'):
return self._run_notebook(script, timeout)
if script.endswith('.py'):
cmd = ['python', script] + args
else:
cmd = [script] + args
# print('@@ cmd={}'.format(cmd))
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
t0, status = time.time(), None
while status is None and time.time() - t0 <= timeout:
time.sleep(1)
status = proc.poll()
if status is None:
dead = self._terminate(proc)
if not dead:
warnings.warn('Unable to terminate process {}'
.format(proc.pid))
return -1
if status != 0:
self._out = proc.stdout.read()
self._err = proc.stderr.read()
return status
def _run_notebook(self, nbfile, timeout):
"""Special case for running a Jupyter Notebook.
See: http://nbconvert.readthedocs.io/en/latest/execute_api.html
"""
status = 0
with open(nbfile) as f:
nb = nbformat.read(f, as_version=4)
ep = ExecutePreprocessor(timeout=timeout + 1)
run_path = os.path.dirname(nbfile)
try:
ep.preprocess(nb, {'metadata': {'path': run_path}})
except Exception as err:
self._err = self.ansi_escape.sub('', str(err))
self._out = 'Error executing the notebook "{}"'.format(nbfile)
status = 1
return status
@property
def output(self):
return self._out
@property
def error(self):
return self._err
@staticmethod
def _terminate(proc):
pid = proc.pid
result = True
os.kill(pid, signal.SIGTERM)
time.sleep(5)
if psutil.pid_exists(pid):
os.kill(pid, signal.SIGKILL)
time.sleep(5)
if psutil.pid_exists(pid):
result = False
return result
def _get_script_path(self):
if os.path.isabs(self._path):
return self._path
return os.path.join(examples_path(), self._path)
[docs]def examples_path():
"""Return path to examples directory.
"""
return os.path.join(os.path.dirname(idaes.__path__[0]),
'examples')
[docs]def find_examples(subdir):
root = os.path.join(examples_path(), subdir)
examples = glob.glob(os.path.join(root, '*_example.py'))
return examples