<commit_before>from __future__ import print_function
from logging import getLogger
logger = getLogger(__name__)

try:
    import thread
except ImportError: # python 3
    import _thread as thread
import threading
import traceback
import socket
import struct
import sys
import os
import atexit
import code
import signal
import errno
try:
    import signalfd
except ImportError:
    signalfd = None
try:
    string = basestring
except NameError: # python 3
    string = str
try:
    InterruptedError = InterruptedError
except NameError: # python <= 3.2
    InterruptedError = OSError
if hasattr(sys, 'setswitchinterval'):
    setinterval = sys.setswitchinterval
    getinterval = sys.getswitchinterval
else:
    setinterval = sys.setcheckinterval
    getinterval = sys.getcheckinterval
PY3 = sys.version_info[0] == 3
VERBOSE = True

try:
    import ctypes
    import ctypes.util
    libpthread_path = ctypes.util.find_library("pthread")
    if not libpthread_path:
        raise ImportError
    libpthread = ctypes.CDLL(libpthread_path)
    if not hasattr(libpthread, "pthread_setname_np"):
        raise ImportError
    _pthread_setname_np = libpthread.pthread_setname_np
    _pthread_setname_np.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
    _pthread_setname_np.restype = ctypes.c_int
    pthread_setname_np = lambda ident, name: _pthread_setname_np(ident, name[:15].encode('utf8'))
except ImportError:
    pthread_setname_np = lambda ident, name: None

SO_PEERCRED = 17

def cry(message):
    """
    Fail-ignorant logging function.
    """
    if VERBOSE:
        try:
            print("Manhole: "+message, file=_STDERR)
        except: #pylint: disable=W0702
            pass

def get_peercred(sock):
    """Gets the (pid, uid, gid) for the client on the given *connected* socket."""

    return struct.unpack('3i', sock.getsockopt(
        socket.SOL_SOCKET, SO_PEERCRED, struct.calcsize('3i')
    ))

class SuspiciousClient(Exception):
    pass

class Manhole(threading.Thread):
    """
    Thread that runs the infamous "Manhole".
    """

    def __init__(self, sigmask):
        super(Manhole, self).__init__()
        self.daemon = True
        self.name = "Manhole"
        self.sigmask = sigmask

    @staticmethod
    def get_socket(factory=None):
        sock = (factory or socket.socket)(socket.AF_UNIX, socket.SOCK_STREAM)
        pid = os.getpid()
        name = "/tmp/manhole-%s" % pid
        if os.path.exists(name):
            os.unlink(name)
        sock.bind(name)
        sock.listen(5)
        cry("Manhole UDS path: "+name)
        return sock, pid

    def run(self):
        if signalfd and self.sigmask:
            signalfd.sigprocmask(signalfd.SIG_BLOCK, self.sigmask)
        pthread_setname_np(self.ident, self.name)

        sock, pid = self.get_socket()
        while True:
            cry("Waiting for new connection (in pid:%s) ..." % pid)
            try:
                client = ManholeConnection(sock.accept()[0], self.sigmask)
                client.start()
                client.join()
            except (InterruptedError, socket.error) as e:
                if e.errno != errno.EINTR:
                    raise
                continue
            finally:
                client = None

class ManholeConnection(threading.Thread):
    def __init__(self, client, sigmask):
        super(ManholeConnection, self).__init__()
        self.daemon = False
        self.client = client
        self.name = "ManholeConnection"
        self.sigmask = sigmask

    def run(self):
        cry('Started ManholeConnection thread. Checking credentials ...')
        if signalfd and self.sigmask:
            signalfd.sigprocmask(signalfd.SIG_BLOCK, self.sigmask)
        pthread_setname_np(self.ident, "Manhole ----")

        pid, _, _ = self.check_credentials(self.client)
        pthread_setname_np(self.ident, "Manhole %s" % pid)
        self.handle(self.client)

    @staticmethod
    def check_credentials(client):
        pid, uid, gid = get_peercred(client)
        euid = os.geteuid()
        client_name = "PID:%s UID:%s GID:%s" % (pid, uid, gid)
        if uid not in (0, euid):
            raise SuspiciousClient(
                "Can't accept client with %s. "
                "It doesn't match the current EUID:%s or ROOT." % (
                    client_name, euid
            ))

        cry("Accepted connection %s from %s" % (client, client_name))
        return pid, uid, gid

    @staticmethod
    def handle(client):
        client.settimeout(None)
        client.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 0)
        client.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 0)
        backup = []
        old_interval = getinterval()
        try:
            try:
                client_fd = client.fileno()
                for mode, names in (
                    ('w', (
                        'stderr',
                        'stdout',
                        '__stderr__',
                        '__stdout__'
                    )),
                    ('r', (
                        'stdin',
                        '__stdin__'
                    ))
                ):
                    for name in names:
                        backup.append((name, getattr(sys, name)))
                        setattr(sys, name, os.fdopen(client_fd, mode, 1 if PY3 else 0))
                run_repl()
                cry("DONE.")
            finally:
                try:
                    setinterval(2147483647) # change the switch/check interval to something ridiculous
                                            # we don't want to have other thread try to write to the
                                            # redirected sys.__std*/sys.std* - it would fail horribly
                    client.close() # close before it's too late. it may already be dead
                    junk = [] # keep the old file objects alive for a bit
                    for name, fh in backup:
                        junk.append(getattr(sys, name))
                        setattr(sys, name, fh)
                    del backup
                    for fh in junk:
                        try:
                            fh.close()
                        except IOError:
                            pass
                        del fh
                    del junk
                finally:
                    setinterval(old_interval)
                    cry("Cleaned up.")
        except Exception:
            cry("ManholeConnection thread failed:")
            cry(traceback.format_exc())

def run_repl():
    dump_stacktraces()
    code.InteractiveConsole({
        'dump_stacktraces': dump_stacktraces,
        'sys': sys,
        'os': os,
        'socket': socket,
        'traceback': traceback,
    }).interact()

def _handle_oneshot(_signum, _frame, _socket=socket._socketobject):
    try:
        sock, pid = Manhole.get_socket(_socket)
        cry("Waiting for new connection (in pid:%s) ..." % pid)
        client, _ = sock.accept()
        ManholeConnection.check_credentials(client)
        ManholeConnection.handle(client)
    except: # pylint: disable=W0702
        # we don't want to let any exception out, it might make the application missbehave
        cry("Manhole oneshot connection failed:")
        cry(traceback.format_exc())
    finally:
        _remove_manhole_uds()

def _remove_manhole_uds():
    name = "/tmp/manhole-%s" % os.getpid()
    if os.path.exists(name):
        os.unlink(name)

_INST_LOCK = thread.allocate_lock()
_STDERR = _INST = _ORIGINAL_OS_FORK = _ORIGINAL_OS_FORKPTY = _SHOULD_RESTART = None

def _patched_fork():
    """Fork a child process."""
    pid = _ORIGINAL_OS_FORK()
    if not pid:
        cry('Fork detected. Reinstalling Manhole.')
        reinstall()
    return pid

def _patched_forkpty():
    """Fork a new process with a new pseudo-terminal as controlling tty."""
    pid, master_fd = _ORIGINAL_OS_FORKPTY()
    if not pid:
        cry('Fork detected. Reinstalling Manhole.')
        reinstall()
    return pid, master_fd

def _patch_os_fork_functions():
    global _ORIGINAL_OS_FORK, _ORIGINAL_OS_FORKPTY #pylint: disable=W0603
    if not _ORIGINAL_OS_FORK:
        _ORIGINAL_OS_FORK, os.fork = os.fork, _patched_fork
    if not _ORIGINAL_OS_FORKPTY:
        _ORIGINAL_OS_FORKPTY, os.forkpty = os.forkpty, _patched_forkpty
    cry("Patched %s and %s." % (_ORIGINAL_OS_FORK, _ORIGINAL_OS_FORKPTY))

def _activate_on_signal(_signum, _frame):
    assert _INST, "Manhole wasn't installed !"
    _INST.start()

ALL_SIGNALS = [
    getattr(signal, sig) for sig in dir(signal)
    if sig.startswith('SIG') and '_' not in sig
]
def install(verbose=True, patch_fork=True, activate_on=None, sigmask=ALL_SIGNALS, oneshot_on=None):
    global _STDERR, _INST, _SHOULD_RESTART, VERBOSE #pylint: disable=W0603
    with _INST_LOCK:
        VERBOSE = verbose
        _STDERR = sys.__stderr__
        if not _INST:
            _INST = Manhole(sigmask)
            if oneshot_on is not None:
                oneshot_on = getattr(signal, 'SIG'+oneshot_on) if isinstance(oneshot_on, string) else oneshot_on
                signal.signal(oneshot_on, _handle_oneshot)

            if activate_on is None:
                if oneshot_on is None:
                    _INST.start()
                    _SHOULD_RESTART = True
            else:
                activate_on = getattr(signal, 'SIG'+activate_on) if isinstance(activate_on, string) else activate_on
                if activate_on == oneshot_on:
                    raise RuntimeError('You cannot do activation of the Manhole thread on the same signal that you want to do oneshot activation !')
                signal.signal(activate_on, _activate_on_signal)
        atexit.register(_remove_manhole_uds)
        if patch_fork:
            if activate_on is None and oneshot_on is None:
                _patch_os_fork_functions()
            else:
                if activate_on:
                    cry("Not patching os.fork and os.forkpty. Activation is done by signal %s" % activate_on)
                elif oneshot_on:
                    cry("Not patching os.fork and os.forkpty. Oneshot activation is done by signal %s" % oneshot_on)

def reinstall():
    global _INST #pylint: disable=W0603
    assert _INST
    with _INST_LOCK:
        if not _INST.is_alive():
            _INST = Manhole(_INST.sigmask)
            if _SHOULD_RESTART:
                _INST.start()

def dump_stacktraces():
    lines = []
    for thread_id, stack in sys._current_frames().items(): #pylint: disable=W0212
        lines.append("\n######### ProcessID=%s, ThreadID=%s #########" % (
            os.getpid(), thread_id
        ))
        for filename, lineno, name, line in traceback.extract_stack(stack):
            lines.append('File: "%s", line %d, in %s' % (filename, lineno, name))
            if line:
                lines.append("  %s" % (line.strip()))
    lines.append("#############################################\n\n")

    print('\n'.join(lines), file=sys.stderr)

if __name__ == '__main__': #pragma: no cover
    from logging import basicConfig, DEBUG
    basicConfig(level=DEBUG)
    install(verbose=True)#, oneshot_on='USR2')
    try:
        import faulthandler
        faulthandler.enable()
    except ImportError:
        pass

    print()

    import time
    from itertools import cycle
    for i, _i in enumerate(cycle([None])):
        if i == 3:
            print()
            print('FORKING ----------------')
            print()
            cpid = os.fork()
            #if cpid:
            #    os.waitpid(pid, 0)
        time.sleep(1)
        #if i == 3:
        #    if cpid:
        #        time.sleep(1000)
        #    else:
        #        run_repl()
        #
<commit_msg>Add complete support for eventlet/gevent.<commit_after>from __future__ import print_function
from logging import getLogger
logger = getLogger(__name__)

import threading
import traceback
import socket
import struct
import sys
import os
import atexit
import code
import signal
import errno
try:
    import signalfd
except ImportError:
    signalfd = None
try:
    string = basestring
except NameError: # python 3
    string = str
try:
    InterruptedError = InterruptedError
except NameError: # python <= 3.2
    InterruptedError = OSError
if hasattr(sys, 'setswitchinterval'):
    setinterval = sys.setswitchinterval
    getinterval = sys.getswitchinterval
else:
    setinterval = sys.setcheckinterval
    getinterval = sys.getcheckinterval

def get_original(qual_name):
    mod, name = qual_name.split('.')
    original = getattr(__import__(mod), name)

    try:
        from gevent.monkey import get_original
        original = get_original(mod, name)
    except ImportError:
        pass

    try:
        from eventlet.patcher import original
        original = getattr(original(mod), name)
    except ImportError:
        pass

    return original
_ORIGINAL_SOCKET = get_original('socket.socket')
_ORIGINAL_FDOPEN = get_original('os.fdopen')
try:
    _ORIGINAL_ALLOCATE_LOCK = get_original('thread.allocate_lock')
except ImportError:
    _ORIGINAL_ALLOCATE_LOCK = get_original('_thread.allocate_lock')  # python3 ...
_ORIGINAL_THREAD = get_original('threading.Thread')

PY3 = sys.version_info[0] == 3
VERBOSE = True

try:
    import ctypes
    import ctypes.util
    libpthread_path = ctypes.util.find_library("pthread")
    if not libpthread_path:
        raise ImportError
    libpthread = ctypes.CDLL(libpthread_path)
    if not hasattr(libpthread, "pthread_setname_np"):
        raise ImportError
    _pthread_setname_np = libpthread.pthread_setname_np
    _pthread_setname_np.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
    _pthread_setname_np.restype = ctypes.c_int
    pthread_setname_np = lambda ident, name: _pthread_setname_np(ident, name[:15].encode('utf8'))
except ImportError:
    pthread_setname_np = lambda ident, name: None

SO_PEERCRED = 17

def cry(message):
    """
    Fail-ignorant logging function.
    """
    if VERBOSE:
        try:
            print("Manhole: "+message, file=_STDERR)
        except: #pylint: disable=W0702
            pass

def get_peercred(sock):
    """Gets the (pid, uid, gid) for the client on the given *connected* socket."""

    return struct.unpack('3i', sock.getsockopt(
        socket.SOL_SOCKET, SO_PEERCRED, struct.calcsize('3i')
    ))

class SuspiciousClient(Exception):
    pass

class Manhole(_ORIGINAL_THREAD):
    """
    Thread that runs the infamous "Manhole".
    """

    def __init__(self, sigmask):
        super(Manhole, self).__init__()
        self.daemon = True
        self.name = "Manhole"
        self.sigmask = sigmask

    @staticmethod
    def get_socket():
        sock = _ORIGINAL_SOCKET(socket.AF_UNIX, socket.SOCK_STREAM)
        pid = os.getpid()
        name = "/tmp/manhole-%s" % pid
        if os.path.exists(name):
            os.unlink(name)
        sock.bind(name)
        sock.listen(5)
        cry("Manhole UDS path: "+name)
        return sock, pid

    def run(self):
        if signalfd and self.sigmask:
            signalfd.sigprocmask(signalfd.SIG_BLOCK, self.sigmask)
        pthread_setname_np(self.ident, self.name)

        sock, pid = self.get_socket()
        while True:
            cry("Waiting for new connection (in pid:%s) ..." % pid)
            try:
                client = ManholeConnection(sock.accept()[0], self.sigmask)
                client.start()
                client.join()
            except (InterruptedError, socket.error) as e:
                if e.errno != errno.EINTR:
                    raise
                continue
            finally:
                client = None

class ManholeConnection(_ORIGINAL_THREAD):
    def __init__(self, client, sigmask):
        super(ManholeConnection, self).__init__()
        self.daemon = False
        self.client = client
        self.name = "ManholeConnection"
        self.sigmask = sigmask

    def run(self):
        cry('Started ManholeConnection thread. Checking credentials ...')
        if signalfd and self.sigmask:
            signalfd.sigprocmask(signalfd.SIG_BLOCK, self.sigmask)
        pthread_setname_np(self.ident, "Manhole ----")

        pid, _, _ = self.check_credentials(self.client)
        pthread_setname_np(self.ident, "Manhole %s" % pid)
        self.handle(self.client)

    @staticmethod
    def check_credentials(client):
        pid, uid, gid = get_peercred(client)
        euid = os.geteuid()
        client_name = "PID:%s UID:%s GID:%s" % (pid, uid, gid)
        if uid not in (0, euid):
            raise SuspiciousClient(
                "Can't accept client with %s. "
                "It doesn't match the current EUID:%s or ROOT." % (
                    client_name, euid
            ))

        cry("Accepted connection %s from %s" % (client, client_name))
        return pid, uid, gid

    @staticmethod
    def handle(client):
        client.settimeout(None)
        client.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 0)
        client.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 0)
        backup = []
        old_interval = getinterval()
        try:
            try:
                client_fd = client.fileno()
                for mode, names in (
                    ('w', (
                        'stderr',
                        'stdout',
                        '__stderr__',
                        '__stdout__'
                    )),
                    ('r', (
                        'stdin',
                        '__stdin__'
                    ))
                ):
                    for name in names:
                        backup.append((name, getattr(sys, name)))
                        setattr(sys, name, _ORIGINAL_FDOPEN(client_fd, mode, 1 if PY3 else 0))
                run_repl()
                cry("DONE.")
            finally:
                try:
                    setinterval(2147483647) # change the switch/check interval to something ridiculous
                                            # we don't want to have other thread try to write to the
                                            # redirected sys.__std*/sys.std* - it would fail horribly
                    client.close() # close before it's too late. it may already be dead
                    junk = [] # keep the old file objects alive for a bit
                    for name, fh in backup:
                        junk.append(getattr(sys, name))
                        setattr(sys, name, fh)
                    del backup
                    for fh in junk:
                        try:
                            fh.close()
                        except IOError:
                            pass
                        del fh
                    del junk
                finally:
                    setinterval(old_interval)
                    cry("Cleaned up.")
        except Exception:
            cry("ManholeConnection thread failed:")
            cry(traceback.format_exc())

def run_repl():
    dump_stacktraces()
    code.InteractiveConsole({
        'dump_stacktraces': dump_stacktraces,
        'sys': sys,
        'os': os,
        'socket': socket,
        'traceback': traceback,
    }).interact()

def _handle_oneshot(_signum, _frame):
    try:
        sock, pid = Manhole.get_socket()
        cry("Waiting for new connection (in pid:%s) ..." % pid)
        client, _ = sock.accept()
        ManholeConnection.check_credentials(client)
        ManholeConnection.handle(client)
    except: # pylint: disable=W0702
        # we don't want to let any exception out, it might make the application missbehave
        cry("Manhole oneshot connection failed:")
        cry(traceback.format_exc())
    finally:
        _remove_manhole_uds()

def _remove_manhole_uds():
    name = "/tmp/manhole-%s" % os.getpid()
    if os.path.exists(name):
        os.unlink(name)

_INST_LOCK = _ORIGINAL_ALLOCATE_LOCK()
_STDERR = _INST = _ORIGINAL_OS_FORK = _ORIGINAL_OS_FORKPTY = _SHOULD_RESTART = None

def _patched_fork():
    """Fork a child process."""
    pid = _ORIGINAL_OS_FORK()
    if not pid:
        cry('Fork detected. Reinstalling Manhole.')
        reinstall()
    return pid

def _patched_forkpty():
    """Fork a new process with a new pseudo-terminal as controlling tty."""
    pid, master_fd = _ORIGINAL_OS_FORKPTY()
    if not pid:
        cry('Fork detected. Reinstalling Manhole.')
        reinstall()
    return pid, master_fd

def _patch_os_fork_functions():
    global _ORIGINAL_OS_FORK, _ORIGINAL_OS_FORKPTY #pylint: disable=W0603
    if not _ORIGINAL_OS_FORK:
        _ORIGINAL_OS_FORK, os.fork = os.fork, _patched_fork
    if not _ORIGINAL_OS_FORKPTY:
        _ORIGINAL_OS_FORKPTY, os.forkpty = os.forkpty, _patched_forkpty
    cry("Patched %s and %s." % (_ORIGINAL_OS_FORK, _ORIGINAL_OS_FORKPTY))

def _activate_on_signal(_signum, _frame):
    assert _INST, "Manhole wasn't installed !"
    _INST.start()

ALL_SIGNALS = [
    getattr(signal, sig) for sig in dir(signal)
    if sig.startswith('SIG') and '_' not in sig
]
def install(verbose=True, patch_fork=True, activate_on=None, sigmask=ALL_SIGNALS, oneshot_on=None):
    global _STDERR, _INST, _SHOULD_RESTART, VERBOSE #pylint: disable=W0603
    with _INST_LOCK:
        VERBOSE = verbose
        _STDERR = sys.__stderr__
        if not _INST:
            _INST = Manhole(sigmask)
            if oneshot_on is not None:
                oneshot_on = getattr(signal, 'SIG'+oneshot_on) if isinstance(oneshot_on, string) else oneshot_on
                signal.signal(oneshot_on, _handle_oneshot)

            if activate_on is None:
                if oneshot_on is None:
                    _INST.start()
                    _SHOULD_RESTART = True
            else:
                activate_on = getattr(signal, 'SIG'+activate_on) if isinstance(activate_on, string) else activate_on
                if activate_on == oneshot_on:
                    raise RuntimeError('You cannot do activation of the Manhole thread on the same signal that you want to do oneshot activation !')
                signal.signal(activate_on, _activate_on_signal)
        atexit.register(_remove_manhole_uds)
        if patch_fork:
            if activate_on is None and oneshot_on is None:
                _patch_os_fork_functions()
            else:
                if activate_on:
                    cry("Not patching os.fork and os.forkpty. Activation is done by signal %s" % activate_on)
                elif oneshot_on:
                    cry("Not patching os.fork and os.forkpty. Oneshot activation is done by signal %s" % oneshot_on)

def reinstall():
    global _INST #pylint: disable=W0603
    assert _INST
    with _INST_LOCK:
        if not _INST.is_alive():
            _INST = Manhole(_INST.sigmask)
            if _SHOULD_RESTART:
                _INST.start()

def dump_stacktraces():
    lines = []
    for thread_id, stack in sys._current_frames().items(): #pylint: disable=W0212
        lines.append("\n######### ProcessID=%s, ThreadID=%s #########" % (
            os.getpid(), thread_id
        ))
        for filename, lineno, name, line in traceback.extract_stack(stack):
            lines.append('File: "%s", line %d, in %s' % (filename, lineno, name))
            if line:
                lines.append("  %s" % (line.strip()))
    lines.append("#############################################\n\n")

    print('\n'.join(lines), file=sys.stderr)

if __name__ == '__main__': #pragma: no cover
    from logging import basicConfig, DEBUG
    basicConfig(level=DEBUG)
    install(verbose=True)#, oneshot_on='USR2')
    try:
        import faulthandler
        faulthandler.enable()
    except ImportError:
        pass

    print()

    import time
    from itertools import cycle
    for i, _i in enumerate(cycle([None])):
        if i == 3:
            print()
            print('FORKING ----------------')
            print()
            cpid = os.fork()
            #if cpid:
            #    os.waitpid(pid, 0)
        time.sleep(1)
        #if i == 3:
        #    if cpid:
        #        time.sleep(1000)
        #    else:
        #        run_repl()
        #
<|endoftext|>


<commit_before>"""Run the E2E workflow.

This script submits Argo workflows to run the E2E tests and waits for
them to finish. It is intended to be invoked by prow jobs.

It requires the workflow to be expressed as a ksonnet app.

The script can take a config file via --config_file.
The --config_file is expected to be a YAML file as follows:

workflows:
  - name: e2e-test
    app_dir: tensorflow/k8s/test/workflows
    component: workflows
    job_types:
      presubmit
    include_dirs:
      tensorflow/*

  - name: lint
    app_dir: kubeflow/kubeflow/testing/workflows
    component: workflows

app_dir is expected to be in the form of
{REPO_OWNER}/{REPO_NAME}/path/to/ksonnet/app

component is the name of the ksonnet component corresponding
to the workflow to launch.

job_types (optional) is an array of strings representing the job types (presubmit/postsubmit)
that should run this workflow.

include_dirs (optional) is an array of strings that specify which directories, if modified,
should run this workflow.

The script expects that the directories
{repos_dir}/{app_dir} exists. Where repos_dir is provided
as a command line argument.
"""

import argparse
import datetime
import fnmatch
import logging
import os
import tempfile
from kubeflow.testing import argo_client
from kubeflow.testing import prow_artifacts
from kubeflow.testing import util
import uuid
import sys
import yaml

# The namespace to launch the Argo workflow in.
def get_namespace(args):
  if args.namespace:
    return args.namespace

  if args.release:
    return "kubeflow-releasing"
  return "kubeflow-test-infra"

class WorkflowComponent(object):
  """Datastructure to represent a ksonnet component to submit a workflow."""

  def __init__(self, name, app_dir, component, job_types, include_dirs, params):
    self.name = name
    self.app_dir = app_dir
    self.component = component
    self.job_types = job_types
    self.include_dirs = include_dirs
    self.params = params

def _get_src_dir():
  return os.path.abspath(os.path.join(__file__, "..",))

def create_started_file(bucket):
  """Create the started file in gcs for gubernator."""
  contents = prow_artifacts.create_started()

  target = os.path.join(prow_artifacts.get_gcs_dir(bucket), "started.json")
  util.upload_to_gcs(contents, target)

def parse_config_file(config_file, root_dir):
  with open(config_file) as hf:
    results = yaml.load(hf)

  components = []
  for i in results["workflows"]:
    components.append(WorkflowComponent(
      i["name"], os.path.join(root_dir, i["app_dir"]), i["component"], i.get("job_types", []),
      i.get("include_dirs", []), i.get("params", {})))
  return components

def generate_env_from_head(args):
  commit = util.run(["git", "rev-parse", "HEAD"], cwd=os.path.join(
    args.repos_dir, os.getenv("REPO_OWNER"), os.getenv("REPO_NAME")))
  pull_base_sha = commit[0:8]
  date_str = datetime.datetime.now().strftime("%Y%m%d")
  build_number = uuid.uuid4().hex[0:4]
  version_tag = "v{0}-{1}".format(date_str, pull_base_sha)
  env_var = {
    "PULL_BASE_SHA": pull_base_sha,
    "BUILD_NUMBER": build_number,
    "VERSION_TAG": version_tag,
  }

  for k in env_var:
    if os.getenv(k):
      continue
    os.environ[k] = env_var.get(k)

def run(args, file_handler): # pylint: disable=too-many-statements,too-many-branches
  # Print ksonnet version
  util.run(["ks", "version"])
  job_type = os.getenv("JOB_TYPE")

  # Compare with master branch and get changed files.
  diff_command = []
  if job_type == "presubmit":
    diff_command = ["git", "diff", "--name-only", "master"]
  else:
    diff_command = ["git", "diff", "--name-only", "HEAD^", "HEAD"]

  changed_files = util.run(diff_command,
    cwd=os.path.join(args.repos_dir, os.getenv("REPO_OWNER"), os.getenv("REPO_NAME"))).splitlines()

  if args.release:
    generate_env_from_head(args)
  workflows = []
  if args.config_file:
    workflows.extend(parse_config_file(args.config_file, args.repos_dir))

  create_started_file(args.bucket)

  util.maybe_activate_service_account()

  util.configure_kubectl(args.project, args.zone, args.cluster)
  util.load_kube_config()

  workflow_names = []
  ui_urls = {}

  for w in workflows:
    # Create the name for the workflow
    # We truncate sha numbers to prevent the workflow name from being too large.
    # Workflow name should not be more than 63 characters because its used
    # as a label on the pods.
    workflow_name = os.getenv("JOB_NAME") + "-" + w.name

    # Skip this workflow if it is scoped to a different job type.
    if w.job_types and not job_type in w.job_types:
      logging.info("Skipping workflow %s because job type %s is not one of "
                   "%s.", w.name, job_type, w.job_types)
      continue

    # If we are scoping this workflow to specific directories, check if any files
    # modified match the specified regex patterns.
    dir_modified = False
    if w.include_dirs:
      for f in changed_files:
        for d in w.include_dirs:
          if fnmatch.fnmatch(f, d):
            dir_modified = True
            logging.info("Triggering workflow %s because %s in dir %s is modified.",
                         w.name, f, d)
            break
        if dir_modified:
          break

    if w.include_dirs and not dir_modified:
      logging.info("Skipping workflow %s because no code modified in %s.",
                   w.name, w.include_dirs)
      continue

    if job_type == "presubmit":
      workflow_name += "-{0}".format(os.getenv("PULL_NUMBER"))
      workflow_name += "-{0}".format(os.getenv("PULL_PULL_SHA")[0:7])

    elif job_type == "postsubmit":
      workflow_name += "-{0}".format(os.getenv("PULL_BASE_SHA")[0:7])

    workflow_name += "-{0}".format(os.getenv("BUILD_NUMBER"))

    salt = uuid.uuid4().hex[0:4]
    # Add some salt. This is mostly a convenience for the case where you
    # are submitting jobs manually for testing/debugging. Since the prow should
    # vend unique build numbers for each job.
    workflow_name += "-{0}".format(salt)

    workflow_names.append(workflow_name)
    # Create a new environment for this run
    env = workflow_name

    util.run(["ks", "env", "add", env], cwd=w.app_dir)

    util.run(["ks", "param", "set", "--env=" + env, w.component,
              "name", workflow_name],
             cwd=w.app_dir)

    # Set the prow environment variables.
    prow_env = []

    names = ["JOB_NAME", "JOB_TYPE", "BUILD_ID", "BUILD_NUMBER",
             "PULL_BASE_SHA", "PULL_NUMBER", "PULL_PULL_SHA", "REPO_OWNER",
             "REPO_NAME"]
    names.sort()
    for v in names:
      if not os.getenv(v):
        continue
      prow_env.append("{0}={1}".format(v, os.getenv(v)))

    util.run(["ks", "param", "set", "--env=" + env, w.component, "prow_env", ",".join(prow_env)],
             cwd=w.app_dir)
    util.run(["ks", "param", "set", "--env=" + env, w.component, "namespace", get_namespace(args)],
             cwd=w.app_dir)
    util.run(["ks", "param", "set", "--env=" + env, w.component, "bucket", args.bucket],
             cwd=w.app_dir)
    if args.release:
      util.run(["ks", "param", "set", "--env=" + env, w.component, "versionTag",
                os.getenv("VERSION_TAG")], cwd=w.app_dir)

    # Set any extra params. We do this in alphabetical order to make it easier to verify in
    # the unittest.
    param_names = w.params.keys()
    param_names.sort()
    for k in param_names:
      util.run(["ks", "param", "set", "--env=" + env, w.component, k, "{0}".format(w.params[k])],
               cwd=w.app_dir)

    # For debugging print out the manifest
    util.run(["ks", "show", env, "-c", w.component], cwd=w.app_dir)
    util.run(["ks", "apply", env, "-c", w.component], cwd=w.app_dir)

    ui_url = ("http://testing-argo.kubeflow.org/workflows/kubeflow-test-infra/{0}"
              "?tab=workflow".format(workflow_name))
    ui_urls[workflow_name] = ui_url
    logging.info("URL for workflow: %s", ui_url)

  success = True
  workflow_phase = {}
  try:
    results = argo_client.wait_for_workflows(get_namespace(args),
                                             workflow_names,
                                             timeout=datetime.timedelta(minutes=180),
                                             status_callback=argo_client.log_status)
    for r in results:
      phase = r.get("status", {}).get("phase")
      name = r.get("metadata", {}).get("name")
      workflow_phase[name] = phase
      if phase != "Succeeded":
        success = False
      logging.info("Workflow %s/%s finished phase: %s", get_namespace(args), name, phase)
  except util.TimeoutError:
    success = False
    logging.exception("Time out waiting for Workflows %s to finish", ",".join(workflow_names))
  except Exception as e:
    # We explicitly log any exceptions so that they will be captured in the
    # build-log.txt that is uploaded to Gubernator.
    logging.exception("Exception occurred: %s", e)
    raise
  finally:
    success = prow_artifacts.finalize_prow_job(args.bucket, success, workflow_phase, ui_urls)

    # Upload logs to GCS. No logs after this point will appear in the
    # file in gcs
    file_handler.flush()
    util.upload_file_to_gcs(
      file_handler.baseFilename,
      os.path.join(prow_artifacts.get_gcs_dir(args.bucket), "build-log.txt"))

  return success

def main(unparsed_args=None):  # pylint: disable=too-many-locals
  logging.getLogger().setLevel(logging.INFO) # pylint: disable=too-many-locals
  # create the top-level parser
  parser = argparse.ArgumentParser(
    description="Submit an Argo workflow to run the E2E tests.")

  parser.add_argument(
    "--project",
    default="",
    type=str,
    help="The project containing the GKE cluster to use to run the workflow.")

  parser.add_argument(
    "--zone",
    default="",
    type=str,
    help="The zone containing the GKE cluster to use to run the workflow.")

  parser.add_argument(
    "--cluster",
    default="",
    type=str,
    help="The GKE cluster to use to run the workflow.")

  parser.add_argument(
    "--bucket",
    default="",
    type=str,
    help="The bucket to use for the Gubernator outputs.")

  parser.add_argument(
    "--config_file",
    default="",
    type=str,
    help="Yaml file containing the config.")

  parser.add_argument(
    "--repos_dir",
    default="",
    type=str,
    help="The directory where the different repos are checked out.")

  parser.add_argument(
    "--release",
    action='store_true',
    default=False,
    help="Whether workflow is for image release")

  parser.add_argument(
    "--namespace",
    default=None,
    type=str,
    help="Optional namespace to use")

  #############################################################################
  # Process the command line arguments.

  # Parse the args
  args = parser.parse_args(args=unparsed_args)

  # Setup a logging file handler. This way we can upload the log outputs
  # to gubernator.
  root_logger = logging.getLogger()

  with tempfile.NamedTemporaryFile(prefix="tmpRunE2eWorkflow", suffix="log") as hf:
    test_log = hf.name

  file_handler = logging.FileHandler(test_log)
  root_logger.addHandler(file_handler)
  # We need to explicitly set the formatter because it will not pick up
  # the BasicConfig.
  formatter = logging.Formatter(fmt=("%(levelname)s|%(asctime)s"
                                     "|%(pathname)s|%(lineno)d| %(message)s"),
                                datefmt="%Y-%m-%dT%H:%M:%S")
  file_handler.setFormatter(formatter)
  logging.info("Logging to %s", test_log)

  return run(args, file_handler)

if __name__ == "__main__":
  logging.basicConfig(level=logging.INFO,
                      format=('%(levelname)s|%(asctime)s'
                              '|%(pathname)s|%(lineno)d| %(message)s'),
                      datefmt='%Y-%m-%dT%H:%M:%S',
                      )
  logging.getLogger().setLevel(logging.INFO)
  final_result = main()
  if not final_result:
    # Exit with a non-zero exit code by to signal failure to prow.
    logging.error("One or more test steps failed exiting with non-zero exit "
                  "code.")
    sys.exit(1)
<commit_msg>Compute postsubmit file diffs by PULL_BASE_SHA (#223)<commit_after>"""Run the E2E workflow.

This script submits Argo workflows to run the E2E tests and waits for
them to finish. It is intended to be invoked by prow jobs.

It requires the workflow to be expressed as a ksonnet app.

The script can take a config file via --config_file.
The --config_file is expected to be a YAML file as follows:

workflows:
  - name: e2e-test
    app_dir: tensorflow/k8s/test/workflows
    component: workflows
    job_types:
      presubmit
    include_dirs:
      tensorflow/*

  - name: lint
    app_dir: kubeflow/kubeflow/testing/workflows
    component: workflows

app_dir is expected to be in the form of
{REPO_OWNER}/{REPO_NAME}/path/to/ksonnet/app

component is the name of the ksonnet component corresponding
to the workflow to launch.

job_types (optional) is an array of strings representing the job types (presubmit/postsubmit)
that should run this workflow.

include_dirs (optional) is an array of strings that specify which directories, if modified,
should run this workflow.

The script expects that the directories
{repos_dir}/{app_dir} exists. Where repos_dir is provided
as a command line argument.
"""

import argparse
import datetime
import fnmatch
import logging
import os
import tempfile
from kubeflow.testing import argo_client
from kubeflow.testing import prow_artifacts
from kubeflow.testing import util
import uuid
import sys
import yaml

# The namespace to launch the Argo workflow in.
def get_namespace(args):
  if args.namespace:
    return args.namespace

  if args.release:
    return "kubeflow-releasing"
  return "kubeflow-test-infra"

class WorkflowComponent(object):
  """Datastructure to represent a ksonnet component to submit a workflow."""

  def __init__(self, name, app_dir, component, job_types, include_dirs, params):
    self.name = name
    self.app_dir = app_dir
    self.component = component
    self.job_types = job_types
    self.include_dirs = include_dirs
    self.params = params

def _get_src_dir():
  return os.path.abspath(os.path.join(__file__, "..",))

def create_started_file(bucket):
  """Create the started file in gcs for gubernator."""
  contents = prow_artifacts.create_started()

  target = os.path.join(prow_artifacts.get_gcs_dir(bucket), "started.json")
  util.upload_to_gcs(contents, target)

def parse_config_file(config_file, root_dir):
  with open(config_file) as hf:
    results = yaml.load(hf)

  components = []
  for i in results["workflows"]:
    components.append(WorkflowComponent(
      i["name"], os.path.join(root_dir, i["app_dir"]), i["component"], i.get("job_types", []),
      i.get("include_dirs", []), i.get("params", {})))
  return components

def generate_env_from_head(args):
  commit = util.run(["git", "rev-parse", "HEAD"], cwd=os.path.join(
    args.repos_dir, os.getenv("REPO_OWNER"), os.getenv("REPO_NAME")))
  pull_base_sha = commit[0:8]
  date_str = datetime.datetime.now().strftime("%Y%m%d")
  build_number = uuid.uuid4().hex[0:4]
  version_tag = "v{0}-{1}".format(date_str, pull_base_sha)
  env_var = {
    "PULL_BASE_SHA": pull_base_sha,
    "BUILD_NUMBER": build_number,
    "VERSION_TAG": version_tag,
  }

  for k in env_var:
    if os.getenv(k):
      continue
    os.environ[k] = env_var.get(k)

def run(args, file_handler): # pylint: disable=too-many-statements,too-many-branches
  # Print ksonnet version
  util.run(["ks", "version"])
  job_type = os.getenv("JOB_TYPE")
  repo_owner = os.getenv("REPO_OWNER")
  repo_name = os.getenv("REPO_NAME")
  pull_base_sha = os.getenv("PULL_BASE_SHA")

  # For presubmit/postsubmit jobs, find the list of files changed by the PR.
  diff_command = []
  if job_type == "presubmit":
    diff_command = ["git", "diff", "--name-only", "master"]
  elif job_type == "postsubmit":
    diff_command = ["git", "diff", "--name-only", pull_base_sha + "^", pull_base_sha]

  changed_files = []
  if job_type == "presubmit" or job_type == "postsubmit":
    changed_files = util.run(diff_command,
      cwd=os.path.join(args.repos_dir, repo_owner, repo_name)).splitlines()

  for f in changed_files:
    logging.info("File %s is modified.", f)

  if args.release:
    generate_env_from_head(args)
  workflows = []
  if args.config_file:
    workflows.extend(parse_config_file(args.config_file, args.repos_dir))

  create_started_file(args.bucket)

  util.maybe_activate_service_account()

  util.configure_kubectl(args.project, args.zone, args.cluster)
  util.load_kube_config()

  workflow_names = []
  ui_urls = {}

  for w in workflows:
    # Create the name for the workflow
    # We truncate sha numbers to prevent the workflow name from being too large.
    # Workflow name should not be more than 63 characters because its used
    # as a label on the pods.
    workflow_name = os.getenv("JOB_NAME") + "-" + w.name

    # Skip this workflow if it is scoped to a different job type.
    if w.job_types and not job_type in w.job_types:
      logging.info("Skipping workflow %s because job type %s is not one of "
                   "%s.", w.name, job_type, w.job_types)
      continue

    # If we are scoping this workflow to specific directories, check if any files
    # modified match the specified regex patterns.
    dir_modified = False
    if w.include_dirs:
      for f in changed_files:
        for d in w.include_dirs:
          if fnmatch.fnmatch(f, d):
            dir_modified = True
            logging.info("Triggering workflow %s because %s in dir %s is modified.",
                         w.name, f, d)
            break
        if dir_modified:
          break

    if w.include_dirs and not dir_modified:
      logging.info("Skipping workflow %s because no code modified in %s.",
                   w.name, w.include_dirs)
      continue

    if job_type == "presubmit":
      workflow_name += "-{0}".format(os.getenv("PULL_NUMBER"))
      workflow_name += "-{0}".format(os.getenv("PULL_PULL_SHA")[0:7])

    elif job_type == "postsubmit":
      workflow_name += "-{0}".format(os.getenv("PULL_BASE_SHA")[0:7])

    workflow_name += "-{0}".format(os.getenv("BUILD_NUMBER"))

    salt = uuid.uuid4().hex[0:4]
    # Add some salt. This is mostly a convenience for the case where you
    # are submitting jobs manually for testing/debugging. Since the prow should
    # vend unique build numbers for each job.
    workflow_name += "-{0}".format(salt)

    workflow_names.append(workflow_name)
    # Create a new environment for this run
    env = workflow_name

    util.run(["ks", "env", "add", env], cwd=w.app_dir)

    util.run(["ks", "param", "set", "--env=" + env, w.component,
              "name", workflow_name],
             cwd=w.app_dir)

    # Set the prow environment variables.
    prow_env = []

    names = ["JOB_NAME", "JOB_TYPE", "BUILD_ID", "BUILD_NUMBER",
             "PULL_BASE_SHA", "PULL_NUMBER", "PULL_PULL_SHA", "REPO_OWNER",
             "REPO_NAME"]
    names.sort()
    for v in names:
      if not os.getenv(v):
        continue
      prow_env.append("{0}={1}".format(v, os.getenv(v)))

    util.run(["ks", "param", "set", "--env=" + env, w.component, "prow_env", ",".join(prow_env)],
             cwd=w.app_dir)
    util.run(["ks", "param", "set", "--env=" + env, w.component, "namespace", get_namespace(args)],
             cwd=w.app_dir)
    util.run(["ks", "param", "set", "--env=" + env, w.component, "bucket", args.bucket],
             cwd=w.app_dir)
    if args.release:
      util.run(["ks", "param", "set", "--env=" + env, w.component, "versionTag",
                os.getenv("VERSION_TAG")], cwd=w.app_dir)

    # Set any extra params. We do this in alphabetical order to make it easier to verify in
    # the unittest.
    param_names = w.params.keys()
    param_names.sort()
    for k in param_names:
      util.run(["ks", "param", "set", "--env=" + env, w.component, k, "{0}".format(w.params[k])],
               cwd=w.app_dir)

    # For debugging print out the manifest
    util.run(["ks", "show", env, "-c", w.component], cwd=w.app_dir)
    util.run(["ks", "apply", env, "-c", w.component], cwd=w.app_dir)

    ui_url = ("http://testing-argo.kubeflow.org/workflows/kubeflow-test-infra/{0}"
              "?tab=workflow".format(workflow_name))
    ui_urls[workflow_name] = ui_url
    logging.info("URL for workflow: %s", ui_url)

  success = True
  workflow_phase = {}
  try:
    results = argo_client.wait_for_workflows(get_namespace(args),
                                             workflow_names,
                                             timeout=datetime.timedelta(minutes=180),
                                             status_callback=argo_client.log_status)
    for r in results:
      phase = r.get("status", {}).get("phase")
      name = r.get("metadata", {}).get("name")
      workflow_phase[name] = phase
      if phase != "Succeeded":
        success = False
      logging.info("Workflow %s/%s finished phase: %s", get_namespace(args), name, phase)
  except util.TimeoutError:
    success = False
    logging.exception("Time out waiting for Workflows %s to finish", ",".join(workflow_names))
  except Exception as e:
    # We explicitly log any exceptions so that they will be captured in the
    # build-log.txt that is uploaded to Gubernator.
    logging.exception("Exception occurred: %s", e)
    raise
  finally:
    success = prow_artifacts.finalize_prow_job(args.bucket, success, workflow_phase, ui_urls)

    # Upload logs to GCS. No logs after this point will appear in the
    # file in gcs
    file_handler.flush()
    util.upload_file_to_gcs(
      file_handler.baseFilename,
      os.path.join(prow_artifacts.get_gcs_dir(args.bucket), "build-log.txt"))

  return success

def main(unparsed_args=None):  # pylint: disable=too-many-locals
  logging.getLogger().setLevel(logging.INFO) # pylint: disable=too-many-locals
  # create the top-level parser
  parser = argparse.ArgumentParser(
    description="Submit an Argo workflow to run the E2E tests.")

  parser.add_argument(
    "--project",
    default="",
    type=str,
    help="The project containing the GKE cluster to use to run the workflow.")

  parser.add_argument(
    "--zone",
    default="",
    type=str,
    help="The zone containing the GKE cluster to use to run the workflow.")

  parser.add_argument(
    "--cluster",
    default="",
    type=str,
    help="The GKE cluster to use to run the workflow.")

  parser.add_argument(
    "--bucket",
    default="",
    type=str,
    help="The bucket to use for the Gubernator outputs.")

  parser.add_argument(
    "--config_file",
    default="",
    type=str,
    help="Yaml file containing the config.")

  parser.add_argument(
    "--repos_dir",
    default="",
    type=str,
    help="The directory where the different repos are checked out.")

  parser.add_argument(
    "--release",
    action='store_true',
    default=False,
    help="Whether workflow is for image release")

  parser.add_argument(
    "--namespace",
    default=None,
    type=str,
    help="Optional namespace to use")

  #############################################################################
  # Process the command line arguments.

  # Parse the args
  args = parser.parse_args(args=unparsed_args)

  # Setup a logging file handler. This way we can upload the log outputs
  # to gubernator.
  root_logger = logging.getLogger()

  with tempfile.NamedTemporaryFile(prefix="tmpRunE2eWorkflow", suffix="log") as hf:
    test_log = hf.name

  file_handler = logging.FileHandler(test_log)
  root_logger.addHandler(file_handler)
  # We need to explicitly set the formatter because it will not pick up
  # the BasicConfig.
  formatter = logging.Formatter(fmt=("%(levelname)s|%(asctime)s"
                                     "|%(pathname)s|%(lineno)d| %(message)s"),
                                datefmt="%Y-%m-%dT%H:%M:%S")
  file_handler.setFormatter(formatter)
  logging.info("Logging to %s", test_log)

  return run(args, file_handler)

if __name__ == "__main__":
  logging.basicConfig(level=logging.INFO,
                      format=('%(levelname)s|%(asctime)s'
                              '|%(pathname)s|%(lineno)d| %(message)s'),
                      datefmt='%Y-%m-%dT%H:%M:%S',
                      )
  logging.getLogger().setLevel(logging.INFO)
  final_result = main()
  if not final_result:
    # Exit with a non-zero exit code by to signal failure to prow.
    logging.error("One or more test steps failed exiting with non-zero exit "
                  "code.")
    sys.exit(1)
<|endoftext|>


<commit_before># ===============================================================================
# Copyright 2013 Jake Ross
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===============================================================================

# ============= enthought library imports =======================
from reportlab.platypus import Paragraph
from traits.api import Bool
from traitsui.api import View
# ============= standard library imports ========================
from reportlab.lib import colors
from reportlab.lib.enums import TA_CENTER
from reportlab.lib.units import inch
# ============= local library imports  ==========================
from pychron.canvas.canvas2D.irradiation_canvas import IrradiationCanvas
# from pychron.entry.level import load_holder_canvas
from pychron.core.pdf.options import BasePDFOptions
from pychron.dvc.meta_repo import irradiation_holder_holes, irradiation_chronology
from pychron.entry.editors.level_editor import load_holder_canvas
from pychron.loading.component_flowable import ComponentFlowable
from pychron.core.pdf.base_table_pdf_writer import BasePDFTableWriter
from pychron.core.pdf.items import Row

class RotatedParagraph(Paragraph):
    rotation = 0

    def draw(self):
        self.canv.saveState()

        self.canv.rotate(self.rotation)
        self.canv.translate(-self.width / 2. - 100, -self.height)
        Paragraph.draw(self)
        self.canv.restoreState()

class IrradiationPDFTableOptions(BasePDFOptions):
    _persistence_name = 'irradiation_pdf_table_options'

    def traits_view(self):
        layout_grp = self._get_layout_group()

        v = View(layout_grp,
                 kind='livemodal',
                 buttons=['OK', 'Cancel'],
                 title='PDF Save Options',
                 resizable=True)
        return v

class IrradiationPDFWriter(BasePDFTableWriter):
    page_break_between_levels = Bool(True)
    show_page_numbers = True
    _options_klass = IrradiationPDFTableOptions

    def _build(self, doc, irrad, *args, **kw):

        self.options.page_number_format = '{} {{page:d}} - {{total:d}}'.format(irrad.name)
        return self._make_levels(irrad)

    def _make_summary(self, irrad):
        fontsize = lambda x, f: '{}'.format(f, x)

        name = irrad.name
        levels = ', '.join(sorted([li.name for li in irrad.levels]))

        chron = irradiation_chronology(name)
        dur = chron.total_duration_seconds
        date = chron.start_date

        # dur = 0
        # if chron:
        #     doses = chron.get_doses()
        #     for pwr, st, en in doses:
        #         dur += (en - st).total_seconds()
        #     _, _, date = chron.get_doses(todatetime=False)[-1]

        dur /= (60 * 60.)
        date = 'Irradiation Date: {}'.format(date)
        dur = 'Irradiation Duration: {:0.1f} hrs'.format(dur)

        name = fontsize(name, 40)
        # levels = fontsize(levels, 28)
        # dur = fontsize(dur, 28)
        txt = '
'.join((name, levels, date, dur))

        klass = Paragraph
        rotation = 0
        if self.options.orientation == 'landscape':
            klass = RotatedParagraph
            rotation = 90

        p = self._new_paragraph(txt,
                                klass=klass,
                                s='Title',
                                textColor=colors.black,
                                alignment=TA_CENTER)
        p.rotation = rotation
        return p

    def _make_levels(self, irrad, progress=None):
        irradname = irrad.name

        p = self._make_summary(irrad)
        flowables = [p, self._page_break()]
        for level in sorted(irrad.levels, key=lambda x: x.name):
            if progress is not None:
                progress.change_message('Making {}{}'.format(irradname, level.name))

            c = self._make_canvas(level)
            fs = self._make_level_table(irrad, level, c)
            if c:
                c = ComponentFlowable(c)
                flowables.append(self._make_table_title(irrad, level))
                flowables.append(c)
                flowables.append(self._page_break())

            flowables.extend(fs)

        return flowables, None

    def _make_level_table(self, irrad, level, c):
        rows = []
        flowables = []
        # flowables.append()
        row = Row()
        row.add_item(span=-1, value=self._make_table_title(irrad, level), fontsize=18)
        rows.append(row)

        row = Row()
        for v in ('', 'Pos.', 'L#', 'Sample', 'Project', 'PI', 'Note'):
            row.add_item(value=self._new_paragraph('{}'.format(v)))
        rows.append(row)

        srows = sorted([self._make_row(pi, c) for pi in level.positions if pi.sample],
                       key=lambda x: x[0])

        rows.extend(srows)

        ts = self._new_style(header_line_idx=0, header_line_width=2)

        ts.add('LINEBELOW', (0, 1), (-1, -1), 1.0, colors.black)

        t = self._new_table(ts, rows)
        t.repeatRows = 2
        t._argW[0] = 0.25 * inch
        t._argW[1] = 0.5 * inch
        t._argW[2] = 0.6 * inch
        t._argW[3] = 1.5 * inch
        t._argW[4] = 1.25 * inch

        flowables.append(t)
        if self.page_break_between_levels:
            flowables.append(self._page_break())
        else:
            flowables.append(self._new_spacer(0, 0.25 * inch))
        return flowables

    def _make_table_title(self, irrad, level):
        t = '{}{} {}'.format(irrad.name, level.name, level.holder)
        p = self._new_paragraph(t, s='Heading1', alignment=TA_CENTER)
        return p

    def _make_row(self, pos, canvas):
        r = Row()
        sample = pos.sample
        project, pi = '', ''
        if sample:
            project = sample.project.name
            pi = sample.project.principal_investigator
            sample = sample.name
            if sample == 'FC-2':
                project, pi = '', ''

        r.add_item(value='[  ]')
        r.add_item(value=pos.position)
        r.add_item(value=pos.identifier)
        r.add_item(value=sample)
        r.add_item(value=project)
        r.add_item(value=pi)
        r.add_item(value='')

        if sample:
            item = canvas.scene.get_item(pos.position)
            item.fill = True

        return r

    def _make_canvas(self, level):
        if level.holder:
            holes = irradiation_holder_holes(level.holder)
            canvas = IrradiationCanvas()
            load_holder_canvas(canvas, holes)
            return canvas

class LabbookPDFWriter(IrradiationPDFWriter):
    def _build(self, doc, irrads, progress=None, *args, **kw):
        flowables = []

        flowables.extend(self._make_title_page(irrads))

        for irrad in irrads:
            fs, _ = self._make_levels(irrad, progress)
            flowables.extend(self._make_summary(irrad))

            flowables.extend(fs)

        return flowables, None

    def _make_title_page(self, irrads):
        start = irrads[0].name
        end = irrads[-1].name
        l1 = 'New Mexico Geochronology Research Laboratory'
        l2 = 'Irradiation Labbook'
        if start != end:
            l3 = '{} to {}'.format(start, end)
        else:
            l3 = start

        t = '
'.join((l1, l2, l3))
        p = self._new_paragraph(t, s='Title')
        return p, self._page_break()

    def _make_summary(self, irrad):
        fontsize = lambda x, f: '{}'.format(f, x)

        name = irrad.name
        levels = ', '.join(sorted([li.name for li in irrad.levels]))

        chron = irradiation_chronology(name)
        dur = chron.total_duration_seconds
        date = chron.start_date

        # dur = 0
        # if chron:
        #     doses = chron.get_doses()
        #     for pwr, st, en in doses:
        #         dur += (en - st).total_seconds()
        #     _, _, date = chron.get_doses(todatetime=False)[-1]

        dur /= (60 * 60.)
        date = 'Irradiation Date: {}'.format(date)
        dur = 'Irradiation Duration: {:0.1f} hrs'.format(dur)

        name = fontsize(name, 40)
        # levels = fontsize(levels, 28)
        # dur = fontsize(dur, 28)
        txt = '
'.join((name, levels, date, dur))
        p = self._new_paragraph(txt,
                                s='Title',
                                textColor=colors.green,
                                alignment=TA_CENTER)

        return p, self._page_break()

# ============= EOF =============================================
<commit_msg>updated irradiation pdf<commit_after># ===============================================================================
# Copyright 2013 Jake Ross
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===============================================================================

# ============= enthought library imports =======================
from reportlab.platypus import Paragraph
from traits.api import Bool
from traitsui.api import View
# ============= standard library imports ========================
from reportlab.lib import colors
from reportlab.lib.enums import TA_CENTER
from reportlab.lib.units import inch
# ============= local library imports  ==========================
from pychron.canvas.canvas2D.irradiation_canvas import IrradiationCanvas
# from pychron.entry.level import load_holder_canvas
from pychron.core.pdf.options import BasePDFOptions
from pychron.dvc.meta_repo import irradiation_holder_holes, irradiation_chronology
from pychron.entry.editors.level_editor import load_holder_canvas
from pychron.loading.component_flowable import ComponentFlowable
from pychron.core.pdf.base_table_pdf_writer import BasePDFTableWriter
from pychron.core.pdf.items import Row

class RotatedParagraph(Paragraph):
    rotation = 0

    def draw(self):
        self.canv.saveState()

        self.canv.rotate(self.rotation)
        self.canv.translate(-self.width / 2. - 100, -self.height)
        Paragraph.draw(self)
        self.canv.restoreState()

class IrradiationPDFTableOptions(BasePDFOptions):
    _persistence_name = 'irradiation_pdf_table_options'

    def traits_view(self):
        layout_grp = self._get_layout_group()

        v = View(layout_grp,
                 kind='livemodal',
                 buttons=['OK', 'Cancel'],
                 title='PDF Save Options',
                 resizable=True)
        return v

class IrradiationPDFWriter(BasePDFTableWriter):
    page_break_between_levels = Bool(True)
    show_page_numbers = True
    _options_klass = IrradiationPDFTableOptions

    def _build(self, doc, irrad, *args, **kw):

        self.options.page_number_format = '{} {{page:d}} - {{total:d}}'.format(irrad.name)
        return self._make_levels(irrad)

    def _make_summary(self, irrad):
        fontsize = lambda x, f: '{}'.format(f, x)

        name = irrad.name
        levels = ', '.join(sorted([li.name for li in irrad.levels]))

        chron = irradiation_chronology(name)
        dur = chron.total_duration_seconds
        date = chron.start_date

        # dur = 0
        # if chron:
        #     doses = chron.get_doses()
        #     for pwr, st, en in doses:
        #         dur += (en - st).total_seconds()
        #     _, _, date = chron.get_doses(todatetime=False)[-1]

        dur /= (60 * 60.)
        date = 'Irradiation Date: {}'.format(date)
        dur = 'Irradiation Duration: {:0.1f} hrs'.format(dur)

        name = fontsize(name, 40)
        # levels = fontsize(levels, 28)
        # dur = fontsize(dur, 28)
        txt = '
'.join((name, levels, date, dur))

        klass = Paragraph
        rotation = 0
        if self.options.orientation == 'landscape':
            klass = RotatedParagraph
            rotation = 90

        p = self._new_paragraph(txt,
                                klass=klass,
                                s='Title',
                                textColor=colors.black,
                                alignment=TA_CENTER)
        p.rotation = rotation
        return p

    def _make_levels(self, irrad, progress=None):
        irradname = irrad.name

        p = self._make_summary(irrad)
        flowables = [p, self._page_break()]
        for level in sorted(irrad.levels, key=lambda x: x.name):
            if progress is not None:
                progress.change_message('Making {}{}'.format(irradname, level.name))

            c = self._make_canvas(level)
            fs = self._make_level_table(irrad, level, c)
            if c:
                c = ComponentFlowable(c)
                flowables.append(self._make_table_title(irrad, level))
                flowables.append(c)
                flowables.append(self._page_break())

            flowables.extend(fs)

        return flowables, None

    def _make_level_table(self, irrad, level, c):
        rows = []
        flowables = []
        # flowables.append()
        row = Row()
        row.add_item(span=-1, value=self._make_table_title(irrad, level), fontsize=18)
        rows.append(row)

        row = Row()
        for v in ('', 'Pos.', 'L#', 'Sample', 'Material', 'Project', 'PI', 'Note'):
            row.add_item(value=self._new_paragraph('{}'.format(v)))
        rows.append(row)

        # srows = []
        # pos = level.positions
        # n = len(pos)
        # # for i,p in enumerate(pos):
        # #     if p.sample:
        # #         srows.append(self._make_row(p, c))
        # #     else:
        # #         # check if the next
        # #         for j in xrange(n-i):
        # #             if

        srows = sorted([self._make_row(pi, c) for pi in level.positions], key=lambda x: x[0])

        rows.extend(srows)

        ts = self._new_style(header_line_idx=0, header_line_width=2)

        ts.add('LINEBELOW', (0, 1), (-1, -1), 1.0, colors.black)

        t = self._new_table(ts, rows)
        t.repeatRows = 2
        t._argW[0] = 0.25 * inch
        t._argW[1] = 0.5 * inch
        t._argW[2] = 0.6 * inch
        t._argW[3] = 1.25 * inch
        t._argW[4] = 0.5 * inch
        t._argW[4] = 1.25 * inch

        flowables.append(t)
        if self.page_break_between_levels:
            flowables.append(self._page_break())
        else:
            flowables.append(self._new_spacer(0, 0.25 * inch))
        return flowables

    def _make_table_title(self, irrad, level):
        t = '{}{} {}'.format(irrad.name, level.name, level.holder)
        p = self._new_paragraph(t, s='Heading1', alignment=TA_CENTER)
        return p

    def _make_row(self, pos, canvas):
        r = Row()
        sample = pos.sample
        project, pi, material = '', '', ''
        if sample:
            if sample.material:
                material = sample.material.name[:15]
            project = sample.project.name
            pi = sample.project.principal_investigator
            sample = sample.name
            if sample == 'FC-2':
                project, pi, material = '', '', ''

        r.add_item(value='[  ]')
        r.add_item(value=pos.position)
        r.add_item(value=pos.identifier or '')
        r.add_item(value=sample or '')
        r.add_item(value=material)
        r.add_item(value=project)
        r.add_item(value=pi)
        r.add_item(value='')

        if sample:
            item = canvas.scene.get_item(pos.position)
            item.fill = True

        return r

    def _make_canvas(self, level):
        if level.holder:
            holes = irradiation_holder_holes(level.holder)
            canvas = IrradiationCanvas()
            load_holder_canvas(canvas, holes)
            return canvas

class LabbookPDFWriter(IrradiationPDFWriter):
    def _build(self, doc, irrads, progress=None, *args, **kw):
        flowables = []

        flowables.extend(self._make_title_page(irrads))

        for irrad in irrads:
            fs, _ = self._make_levels(irrad, progress)
            flowables.extend(self._make_summary(irrad))

            flowables.extend(fs)

        return flowables, None

    def _make_title_page(self, irrads):
        start = irrads[0].name
        end = irrads[-1].name
        l1 = 'New Mexico Geochronology Research Laboratory'
        l2 = 'Irradiation Labbook'
        if start != end:
            l3 = '{} to {}'.format(start, end)
        else:
            l3 = start

        t = '
'.join((l1, l2, l3))
        p = self._new_paragraph(t, s='Title')
        return p, self._page_break()

    def _make_summary(self, irrad):
        fontsize = lambda x, f: '{}'.format(f, x)

        name = irrad.name
        levels = ', '.join(sorted([li.name for li in irrad.levels]))

        chron = irradiation_chronology(name)
        dur = chron.total_duration_seconds
        date = chron.start_date

        # dur = 0
        # if chron:
        #     doses = chron.get_doses()
        #     for pwr, st, en in doses:
        #         dur += (en - st).total_seconds()
        #     _, _, date = chron.get_doses(todatetime=False)[-1]

        dur /= (60 * 60.)
        date = 'Irradiation Date: {}'.format(date)
        dur = 'Irradiation Duration: {:0.1f} hrs'.format(dur)

        name = fontsize(name, 40)
        # levels = fontsize(levels, 28)
        # dur = fontsize(dur, 28)
        txt = '
'.join((name, levels, date, dur))
        p = self._new_paragraph(txt,
                                s='Title',
                                textColor=colors.green,
                                alignment=TA_CENTER)

        return p, self._page_break()

# ============= EOF =============================================
<|endoftext|>


<commit_before># ===============================================================================
# Copyright 2015 Jake Ross
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===============================================================================

from __future__ import absolute_import
from envisage.ui.tasks.task_factory import TaskFactory
from traits.api import List

from pychron.envisage.tasks.base_task_plugin import BaseTaskPlugin
from pychron.furnace.ifurnace_manager import IFurnaceManager
from pychron.furnace.tasks.preferences import NMGRLFurnacePreferencesPane
from pychron.furnace.tasks.preferences import LDEOFurnacePreferencesPane
from pychron.furnace.tasks.task import FurnaceTask, LDEOFurnaceTask

class BaseFurnacePlugin(BaseTaskPlugin):
    id = 'pychron.furnace.base.plugin'
    managers = List(contributes_to='pychron.hardware.managers')
    klass = None

    def _manager_factory(self):
        factory = __import__(self.klass[0], fromlist=[self.klass[1]])
        m = getattr(factory, self.klass[1])()

        # m.bootstrap()
        m.name = self.name
        m.plugin_id = self.id
        # m.bind_preferences(self.id)

        return m

    def _task_factory(self):
        return FurnaceTask(manager=self._get_manager(),
                           application=self.application)

    def _get_manager(self):
        return self.application.get_service(IFurnaceManager, 'name=="{}"'.format(self.name))

    def _managers_default(self):
        d = []

        if self.klass is not None:
            d = [dict(name=self.name,
                      plugin_name=self.name,
                      manager=self.application.get_service(IFurnaceManager, 'name=="{}"'.format(self.name)))]

        return d

    def _service_offers_default(self):
        so = self.service_offer_factory(protocol=IFurnaceManager,
                                        factory=self._manager_factory)
        return [so]

    def _tasks_default(self):
        return [TaskFactory(name='Furnace',
                            task_group='hardware',
                            factory=self._task_factory,
                            protocol=FurnaceTask)]

class NMGRLFurnacePlugin(BaseFurnacePlugin):
    name = 'NMGRLFurnace'
    id = 'pychron.furnace.nmgrl.plugin'

    klass = ('pychron.furnace.furnace_manager', 'NMGRLFurnaceManager')

    panes = List(contributes_to='pychron.experiment.dock_pane_factories')
    activations = List(contributes_to='pychron.experiment.activations')
    deactivations = List(contributes_to='pychron.experiment.deactivations')

    def _help_tips_default(self):
        return ['NMGRLFurnace hardware was designed by William McIntosh. Firmware and software was designed by Jake '
                'Ross']

    def _deactivations_default(self):
        application = self.application

        def func():
            manager = application.get_service(IFurnaceManager)
            if manager:
                for window in application.windows:
                    if 'furnace' in window.active_task.id:
                        break
                else:
                    manager.stop_update()

        return [func]

    def _activations_default(self):
        man = self._get_manager()
        return [man.start_update]

    def _panes_default(self):
        def f():
            from pychron.furnace.tasks.panes import ExperimentFurnacePane
            manager = self._get_manager()
            fpane = ExperimentFurnacePane(model=manager)
            return fpane

        return [f]

    def test_furnace_api(self):
        man = self._get_manager()
        return man.test_furnace_api()

    def test_furnace_cam(self):
        man = self._get_manager()
        return man.test_furnace_cam()

    def _preferences_panes_default(self):
        return [NMGRLFurnacePreferencesPane]

class LDEOFurnacePlugin(BaseFurnacePlugin):
    name = 'LDEOFurnace'
    id = 'pychron.furnace.ldeo.plugin'

    klass = ('pychron.furnace.furnace_manager', 'LDEOFurnaceManager')

    panes = List(contributes_to='pychron.experiment.dock_pane_factories')
    activations = List(contributes_to='pychron.experiment.activations')
    deactivations = List(contributes_to='pychron.experiment.deactivations')

    def _task_factory(self):
        return LDEOFurnaceTask(manager=self._get_manager(),
                           application=self.application)

    def _help_tips_default(self):
        return ['LDEOFurnace hardware was designed by Stephen Cox. Firmware and software was designed by Jake '
                'Ross for NMGRL with modifications by Stephen Cox for LDEO']

    def _deactivations_default(self):
        application = self.application

        def func():
            manager = application.get_service(IFurnaceManager)
            if manager:
                for window in application.windows:
                    if 'furnace' in window.active_task.id:
                        break
                else:
                    manager.stop_update()

        return [func]

    def _activations_default(self):
        man = self._get_manager()
        return [man.start_update]

    def _panes_default(self):
        def f():
            from pychron.furnace.tasks.panes import ExperimentFurnacePane
            manager = self._get_manager()
            fpane = ExperimentFurnacePane(model=manager)
            return fpane

        return [f]

    def test_furnace_api(self):
        man = self._get_manager()
        return man.test_furnace_api()

    def test_furnace_cam(self):
        man = self._get_manager()
        return man.test_furnace_cam()

    def _preferences_panes_default(self):
        return [LDEOFurnacePreferencesPane]

    def _tasks_default(self):
        return [TaskFactory(name='Furnace',
                            task_group='hardware',
                            factory=self._task_factory,
                            protocol=LDEOFurnaceTask)]

# ============= EOF =============================================
<commit_msg>Pref pane addition<commit_after># ===============================================================================
# Copyright 2015 Jake Ross
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===============================================================================

from __future__ import absolute_import
from envisage.ui.tasks.task_factory import TaskFactory
from traits.api import List

from pychron.envisage.tasks.base_task_plugin import BaseTaskPlugin
from pychron.furnace.ifurnace_manager import IFurnaceManager
from pychron.furnace.tasks.preferences import NMGRLFurnacePreferencesPane
from pychron.furnace.tasks.preferences import LDEOFurnacePreferencesPane, LDEOFurnaceControlPreferencesPane
from pychron.furnace.tasks.task import FurnaceTask, LDEOFurnaceTask

class BaseFurnacePlugin(BaseTaskPlugin):
    id = 'pychron.furnace.base.plugin'
    managers = List(contributes_to='pychron.hardware.managers')
    klass = None

    def _manager_factory(self):
        factory = __import__(self.klass[0], fromlist=[self.klass[1]])
        m = getattr(factory, self.klass[1])()

        # m.bootstrap()
        m.name = self.name
        m.plugin_id = self.id
        # m.bind_preferences(self.id)

        return m

    def _task_factory(self):
        return FurnaceTask(manager=self._get_manager(),
                           application=self.application)

    def _get_manager(self):
        return self.application.get_service(IFurnaceManager, 'name=="{}"'.format(self.name))

    def _managers_default(self):
        d = []

        if self.klass is not None:
            d = [dict(name=self.name,
                      plugin_name=self.name,
                      manager=self.application.get_service(IFurnaceManager, 'name=="{}"'.format(self.name)))]

        return d

    def _service_offers_default(self):
        so = self.service_offer_factory(protocol=IFurnaceManager,
                                        factory=self._manager_factory)
        return [so]

    def _tasks_default(self):
        return [TaskFactory(name='Furnace',
                            task_group='hardware',
                            factory=self._task_factory,
                            protocol=FurnaceTask)]

class NMGRLFurnacePlugin(BaseFurnacePlugin):
    name = 'NMGRLFurnace'
    id = 'pychron.furnace.nmgrl.plugin'

    klass = ('pychron.furnace.furnace_manager', 'NMGRLFurnaceManager')

    panes = List(contributes_to='pychron.experiment.dock_pane_factories')
    activations = List(contributes_to='pychron.experiment.activations')
    deactivations = List(contributes_to='pychron.experiment.deactivations')

    def _help_tips_default(self):
        return ['NMGRLFurnace hardware was designed by William McIntosh. Firmware and software was designed by Jake '
                'Ross']

    def _deactivations_default(self):
        application = self.application

        def func():
            manager = application.get_service(IFurnaceManager)
            if manager:
                for window in application.windows:
                    if 'furnace' in window.active_task.id:
                        break
                else:
                    manager.stop_update()

        return [func]

    def _activations_default(self):
        man = self._get_manager()
        return [man.start_update]

    def _panes_default(self):
        def f():
            from pychron.furnace.tasks.panes import ExperimentFurnacePane
            manager = self._get_manager()
            fpane = ExperimentFurnacePane(model=manager)
            return fpane

        return [f]

    def test_furnace_api(self):
        man = self._get_manager()
        return man.test_furnace_api()

    def test_furnace_cam(self):
        man = self._get_manager()
        return man.test_furnace_cam()

    def _preferences_panes_default(self):
        return [NMGRLFurnacePreferencesPane]

class LDEOFurnacePlugin(BaseFurnacePlugin):
    name = 'LDEOFurnace'
    id = 'pychron.furnace.ldeo.plugin'

    klass = ('pychron.furnace.furnace_manager', 'LDEOFurnaceManager')

    panes = List(contributes_to='pychron.experiment.dock_pane_factories')
    activations = List(contributes_to='pychron.experiment.activations')
    deactivations = List(contributes_to='pychron.experiment.deactivations')

    def _task_factory(self):
        return LDEOFurnaceTask(manager=self._get_manager(),
                           application=self.application)

    def _help_tips_default(self):
        return ['LDEOFurnace hardware was designed by Stephen Cox. Firmware and software was designed by Jake '
                'Ross for NMGRL with modifications by Stephen Cox for LDEO']

    def _deactivations_default(self):
        application = self.application

        def func():
            manager = application.get_service(IFurnaceManager)
            if manager:
                for window in application.windows:
                    if 'furnace' in window.active_task.id:
                        break
                else:
                    manager.stop_update()

        return [func]

    def _activations_default(self):
        man = self._get_manager()
        return [man.start_update]

    def _panes_default(self):
        def f():
            from pychron.furnace.tasks.panes import ExperimentFurnacePane
            manager = self._get_manager()
            fpane = ExperimentFurnacePane(model=manager)
            return fpane

        return [f]

    def test_furnace_api(self):
        man = self._get_manager()
        return man.test_furnace_api()

    def test_furnace_cam(self):
        man = self._get_manager()
        return man.test_furnace_cam()

    def _preferences_panes_default(self):
        return [LDEOFurnacePreferencesPane, LDEOFurnaceControlPreferencesPane]

    def _tasks_default(self):
        return [TaskFactory(name='Furnace',
                            task_group='hardware',
                            factory=self._task_factory,
                            protocol=LDEOFurnaceTask)]

# ============= EOF =============================================
<|endoftext|>


<commit_before>
from unittest import TestCase
import pycqed as pq
import os

# from qcodes import Instrument

from pycqed.instrument_drivers.physical_instruments._controlbox \
    import qasm_compiler as qasm_compiler

from pycqed.utilities.general import mopen
from pycqed.measurement.waveform_control_CC import \
    single_qubit_qasm_seqs as sq_qasm

# from pycqed.instrument_drivers.meta_instrument.qubit_objects.qubit_object \
#     import Transmon
from pycqed.instrument_drivers.physical_instruments._controlbox \
    import Assembler

from pycqed.measurement.waveform_control_CC import qasm_to_asm as qta

class Test_Parallel_QASM_Compiler(TestCase):

    @classmethod
    def setUpClass(self):
        print("hello world.")

    def setCompiler(self, qasm_file_name=None):
        test_file_dir = os.path.join(
            pq.__path__[0], 'tests', 'test_data', "qasm_test")
        if qasm_file_name != None:
            self.qasm_file_name = os.path.join(test_file_dir,
                                                qasm_file_name)
        else:
            self.qasm_file_name = os.path.join(test_file_dir,
                                                "test_qasm.qasm")

        self.compiler = qasm_compiler.QASM_QuMIS_Compiler(self.qasm_file_name)

    def test_read_file(self):
        qasm_file_name = "Nonexisting.qasm"
        self.setCompiler(qasm_file_name)
        with self.assertRaises(OSError):
            self.compiler.read_file()

    def test_remove_comment(self):
        self.assertTrue(True)

# tpac = Test_Parallel_QASM_Compiler()
# tpac.test_remove_comment()<commit_msg>removed the test file for parallel qasm compiler.<commit_after><|endoftext|>


<commit_before>"""
Functions for computing likelihoods of linear mixed models
"""

from math import log, pi

import numpy as np
from scipy.sparse import bsr_matrix, issparse
from scipy.linalg import pinv, inv
from scipy import matrix
np.seterr(invalid='ignore')

l2pi = log(2 * pi)

def logdet(M):
    """ Returns the (positive) log determinant of a matrix. """
    sign, logdet = np.linalg.slogdet(M.todense() if issparse(M) else M)
    return logdet

def makeR(y, X, V=None, Vinv=None):
    """ Makes the R matrix commonly found in REML estimation """
    if V is None and Vinv is None:
        raise ValueError('Variance matrix not specified')
    elif Vinv is None and V is not None:
        Vinv = makeVinv(V)
    return y - X * pinv(X.transpose() * Vinv * X) * X.transpose() * Vinv * y

def makeP(X, Vinv):
    """ Makes the P matrix commonly found in mixed model estimation """
    return Vinv - Vinv * X * inv(X.T * Vinv * X) * X.T * Vinv

def makeVinv(V):
    if issparse(V):
        V = V.todense()
    return bsr_matrix(inv(V))

def full_loglikelihood(y, V, X, beta, Vinv=None):
    """
    Returns the full loglikelihood of a mixed model

    Ref: SAS documentation for PROC MIXED
    """
    if Vinv is None:
        Vinv = makeVinv(V)
    n = X.shape[0]
    fixefresids = y - X * beta
    llik = -0.5 * ( n * l2pi + logdet(V) + fixefresids.T * Vinv * fixefresids)
    return matrix.item(llik)

def restricted_loglikelihood(y, V, X, P=None, Vinv=None):
    """
    Returns the restricted loglikelihood for mixed model variance component
    estimation.

    References:

    Harville. 'Maximum Likelihood Approaches to Variance Component Estimation
    and to Related Problems' Journal of the American Statistical Association.
    (1977) (72):258

    Lange, Westlake, & Spence. 'Extensions to pedigree analysis III. Variance
    components by the scoring method.' Ann Hum Genet. (1976). 39:4,485-491
    DOI: 10.1111/j.1469-1809.1976.tb00156.x

    SAS documentation for PROC MIXED
    """
    if Vinv is None:
        Vinv = makeVinv(V)
    if P is None:
        P = makeP(X, Vinv=Vinv)
    n = X.shape[0]
    rank = np.linalg.matrix_rank(X)
    llik_restricted = -0.5 * (logdet(V.todense())
                              + logdet(X.transpose() * Vinv * X)
                              + y.T * P * y
                              + (n - rank) * l2pi)
    return matrix.item(llik_restricted)

def reml_gradient(y, X, V, ranefs, P=None, Vinv=None):
    if Vinv is None:
        Vinv = makeVinv(V)
    if P is None:
        P = makeP(X, Vinv)
    nabla = [dREML_dsigma(y, rf.Z, rf.G, P) for rf in ranefs]
    return np.array(nabla)

def dREML_dsigma(y, Z, G, P):
    "The REML derivative of V with regard to sigma"
    PZGZt = P * Z * G * Z.T
    dl_dsig = -.5 * np.trace(PZGZt) + .5 * (y.T * PZGZt * P * y)
    return matrix.item(dl_dsig)

def reml_hessian_element(y, P, dV_dsigma_a, dV_dsigma_b):
    common_term = P * dV_dsigma_a * P * dV_dsigma_b
    a = .5 * np.trace(common_term)
    b = y.T * common_term * P * y
    return matrix.item(a - b)

def reml_hessian(y, X, V, ranefs, P=None, Vinv=None):
    if Vinv is None:
        Vinv = makeVinv(V)
    if P is None:
        P = makeP(X, Vinv)

    n_ranefs = len(ranefs)
    mat = np.zeros((n_ranefs, n_ranefs))

    for i, ranef_a in enumerate(ranefs):
        dV_dsigma_a = ranef_a.V_i
        for j, ranef_b in enumerate(ranefs):

            if j < i:
                # Already set when we did the other side of the matrix
                continue

            dV_dsigma_b = ranef_b.V_i
            element = reml_hessian_element(y, P, dV_dsigma_a, dV_dsigma_b)

            mat[i, j] = element
            mat[j, i] = element

    return np.array(mat, dtype=np.float128)

def reml_observed_information_matrix(y, X, V, ranefs, P=None, Vinv=None):
    return -reml_hessian(y, X, V, ranefs, P, Vinv)

def reml_fisher_element(P, dV_dsigma_a, dV_dsigma_b):
    return .5 * np.trace(P * dV_dsigma_a * P * dV_dsigma_b)

def reml_fisher_information_matrix(y, X, V, ranefs, P=None, Vinv=None):
    if Vinv is None:
        Vinv = makeVinv(V)
    if P is None:
        P = makeP(X, Vinv)

    mat = np.zeros((len(ranefs), len(ranefs)), dtype=np.float128)

    for i, ranef_a in enumerate(ranefs):
        dV_dsigma_a = ranef_a.V_i

        for j, ranef_b in enumerate(ranefs):
            if j < i:
                # Already set when we did the other side of the matrix
                continue

            element = reml_fisher_element(P, dV_dsigma_a, ranef_b.V_i)

            mat[i, j] = element
            mat[j, i] = element

    return mat

def reml_average_information_element(y, P, dV_dsigma_a, dV_dsigma_b):
    return .5 * y.T * P * dV_dsigma_a * P * dV_dsigma_b * P * y

def reml_average_information_matrix(y, X, V, ranefs, P=None, Vinv=None):
    if Vinv is None:
        Vinv = makeVinv(V)
    if P is None:
        P = makeP(X, Vinv)
    mat = np.zeros((len(ranefs), len(ranefs)), dtype=np.float128)
    for i, ranef_a in enumerate(ranefs):
        dV_dsigma_a = ranef_a.V_i

        for j, ranef_b in enumerate(ranefs):
            if j < i:
                # Already set when we did the other side of the matrix
                continue
            dV_dsigma_b = ranef_b.V_i
            element = reml_average_information_element(y, P,
                                                       dV_dsigma_a,
                                                       dV_dsigma_b)
            mat[i, j] = element
            mat[j, i] = element

    return mat
<commit_msg>float128 isnt worth the extra precision if np.linalg doesnt support it<commit_after>"""
Functions for computing likelihoods of linear mixed models
"""

from math import log, pi

import numpy as np
from scipy.sparse import bsr_matrix, issparse
from scipy.linalg import pinv, inv
from scipy import matrix
# np.seterr(invalid='ignore')

l2pi = log(2 * pi)

def logdet(M):
    """ Returns the (positive) log determinant of a matrix. """
    sign, logdet = np.linalg.slogdet(M.todense() if issparse(M) else M)
    return logdet

def makeR(y, X, V=None, Vinv=None):
    """ Makes the R matrix commonly found in REML estimation """
    if V is None and Vinv is None:
        raise ValueError('Variance matrix not specified')
    elif Vinv is None and V is not None:
        Vinv = makeVinv(V)
    return y - X * pinv(X.transpose() * Vinv * X) * X.transpose() * Vinv * y

def makeP(X, Vinv):
    """ Makes the P matrix commonly found in mixed model estimation """
    return Vinv - Vinv * X * inv(X.T * Vinv * X) * X.T * Vinv

def makeVinv(V):
    if issparse(V):
        V = V.todense()
    return bsr_matrix(inv(V))

def full_loglikelihood(y, V, X, beta, Vinv=None):
    """
    Returns the full loglikelihood of a mixed model

    Ref: SAS documentation for PROC MIXED
    """
    if Vinv is None:
        Vinv = makeVinv(V)
    n = X.shape[0]
    fixefresids = y - X * beta
    llik = -0.5 * ( n * l2pi + logdet(V) + fixefresids.T * Vinv * fixefresids)
    return matrix.item(llik)

def restricted_loglikelihood(y, V, X, P=None, Vinv=None):
    """
    Returns the restricted loglikelihood for mixed model variance component
    estimation.

    References:

    Harville. 'Maximum Likelihood Approaches to Variance Component Estimation
    and to Related Problems' Journal of the American Statistical Association.
    (1977) (72):258

    Lange, Westlake, & Spence. 'Extensions to pedigree analysis III. Variance
    components by the scoring method.' Ann Hum Genet. (1976). 39:4,485-491
    DOI: 10.1111/j.1469-1809.1976.tb00156.x

    SAS documentation for PROC MIXED
    """
    if Vinv is None:
        Vinv = makeVinv(V)
    if P is None:
        P = makeP(X, Vinv=Vinv)
    n = X.shape[0]
    rank = np.linalg.matrix_rank(X)
    llik_restricted = -0.5 * (logdet(V.todense())
                              + logdet(X.transpose() * Vinv * X)
                              + y.T * P * y
                              + (n - rank) * l2pi)
    return matrix.item(llik_restricted)

def reml_gradient(y, X, V, ranefs, P=None, Vinv=None):
    if Vinv is None:
        Vinv = makeVinv(V)
    if P is None:
        P = makeP(X, Vinv)
    nabla = [dREML_dsigma(y, rf.Z, rf.G, P) for rf in ranefs]
    return np.array(nabla)

def dREML_dsigma(y, Z, G, P):
    "The REML derivative of V with regard to sigma"
    PZGZt = P * Z * G * Z.T
    dl_dsig = -.5 * np.trace(PZGZt) + .5 * (y.T * PZGZt * P * y)
    return matrix.item(dl_dsig)

def reml_hessian_element(y, P, dV_dsigma_a, dV_dsigma_b):
    common_term = P * dV_dsigma_a * P * dV_dsigma_b
    a = .5 * np.trace(common_term)
    b = y.T * common_term * P * y
    return matrix.item(a - b)

def reml_hessian(y, X, V, ranefs, P=None, Vinv=None):
    if Vinv is None:
        Vinv = makeVinv(V)
    if P is None:
        P = makeP(X, Vinv)

    n_ranefs = len(ranefs)
    mat = np.zeros((n_ranefs, n_ranefs))

    for i, ranef_a in enumerate(ranefs):
        dV_dsigma_a = ranef_a.V_i
        for j, ranef_b in enumerate(ranefs):

            if j < i:
                # Already set when we did the other side of the matrix
                continue

            dV_dsigma_b = ranef_b.V_i
            element = reml_hessian_element(y, P, dV_dsigma_a, dV_dsigma_b)

            mat[i, j] = element
            mat[j, i] = element

    return np.array(mat)

def reml_observed_information_matrix(y, X, V, ranefs, P=None, Vinv=None):
    return -reml_hessian(y, X, V, ranefs, P, Vinv)

def reml_fisher_element(P, dV_dsigma_a, dV_dsigma_b):
    return .5 * np.trace(P * dV_dsigma_a * P * dV_dsigma_b)

def reml_fisher_information_matrix(y, X, V, ranefs, P=None, Vinv=None):
    if Vinv is None:
        Vinv = makeVinv(V)
    if P is None:
        P = makeP(X, Vinv)

    mat = np.zeros((len(ranefs), len(ranefs)))

    for i, ranef_a in enumerate(ranefs):
        dV_dsigma_a = ranef_a.V_i

        for j, ranef_b in enumerate(ranefs):
            if j < i:
                # Already set when we did the other side of the matrix
                continue

            element = reml_fisher_element(P, dV_dsigma_a, ranef_b.V_i)

            mat[i, j] = element
            mat[j, i] = element

    return mat

def reml_average_information_element(y, P, dV_dsigma_a, dV_dsigma_b):
    return .5 * y.T * P * dV_dsigma_a * P * dV_dsigma_b * P * y

def reml_average_information_matrix(y, X, V, ranefs, P=None, Vinv=None):
    if Vinv is None:
        Vinv = makeVinv(V)
    if P is None:
        P = makeP(X, Vinv)
    mat = np.zeros((len(ranefs), len(ranefs)))
    for i, ranef_a in enumerate(ranefs):
        dV_dsigma_a = ranef_a.V_i

        for j, ranef_b in enumerate(ranefs):
            if j < i:
                # Already set when we did the other side of the matrix
                continue
            dV_dsigma_b = ranef_b.V_i
            element = reml_average_information_element(y, P,
                                                       dV_dsigma_a,
                                                       dV_dsigma_b)
            mat[i, j] = element
            mat[j, i] = element

    return mat
<|endoftext|>


<commit_before>__version__ = '0.2.0'
<commit_msg>Bump version to 0.3.0<commit_after>__version__ = '0.3.0'
<|endoftext|>