<commit_before>#!/usr/bin/env python

import argparse
import os
import subprocess
import sys
import threading
from collections import OrderedDict

import yaml
import yodl

debug = False

class DockerCompose:
    def __init__(self, compose, project, compose_base_dir, requested_services):
        self.project = project
        self.compose_base_dir = compose_base_dir
        self.services = self.merge_services(compose.get('services', {}))
        self.networks = compose.get('networks', {})
        self.volumes = compose.get('volumes', {})
        self.filtered_services = filter(lambda service: not requested_services or service in requested_services, self.services)

    def project_prefix(self, value):
        return '{}_{}'.format(self.project, value)

    def merge_services(self, services):
        result = OrderedDict()

        for service in services:
            service_config = services[service]
            result[service] = service_config

            if 'extends' in service_config:
                extended_config = service_config['extends']
                extended_service = extended_config['service']

                del result[service]['extends']

                if 'file' in extended_config:
                    extended_service_data = self.merge_services(
                        yaml.load(open(self.compose_base_dir + extended_config['file'], 'r'), yodl.OrderedDictYAMLLoader)['services']
                    )[extended_service]
                else:
                    extended_service_data = result[extended_service]

                for k, v in extended_service_data.items():
                    if k not in result[service]:
                        result[service][k] = v

        return result

    @staticmethod
    def call(cmd, ignore_return_code=False):
        print('Running: \n' + cmd + '\n')
        if not debug:
            ps = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
            returncode = ps.wait()
            stdout = ps.communicate()[0]
            if returncode != 0 and not ignore_return_code:
                print >> sys.stderr, ('Error: command "{}" failed: {}'.format(cmd, stdout))
                sys.exit(returncode)
            else:
                return stdout

    def is_service_exists(self, service):
        return self.call('docker service ls -f name={} | tail -n +2'.format(self.project_prefix(service)))

    def is_external_network(self, network):
        return isinstance(self.networks[network], dict) and 'external' in self.networks[network]

    def up(self):
        for network in self.networks:
            if not self.is_external_network(network):
                cmd = '[ "`docker network ls -f name={0} | tail -n +2`" != "" ] || docker network create --driver overlay --opt encrypted {0}' \
                    .format(self.project_prefix(network))
                self.call(cmd)

        for volume in self.volumes:
            cmd = '[ "`docker volume ls -f name={0} | tail -n +2`" != "" ] || docker volume create --name {0}'.format(self.project_prefix(volume))
            self.call(cmd)

        services_to_start = []

        for service in self.filtered_services:
            if self.is_service_exists(service):
                services_to_start.append(service)
                continue

            service_config = self.services[service]
            cmd = ['docker service create --with-registry-auth \\\n --name', self.project_prefix(service), '\\\n']

            service_image = []
            service_command = []

            for parameter in service_config:
                value = service_config[parameter]

                def restart():
                    cmd.extend(['--restart-condition', {'always': 'any'}[value], '\\\n'])

                def logging():
                    cmd.extend(['--log-driver', value['driver'], '\\\n'])
                    log_opts = value['options']
                    for k, v in log_opts.items():
                        cmd.extend(['--log-opt', '{}={}'.format(k, v), '\\\n'])

                def mem_limit():
                    cmd.extend(['--limit-memory', value, '\\\n'])

                def image():
                    service_image.append(value)

                def command():
                    service_command.extend(value.split(' '))

                def expose():
                    pass  # unsupported

                def container_name():
                    pass  # unsupported

                def hostname():
                    pass  # unsupported; waiting for https://github.com/docker/docker/issues/24877

                def extra_hosts():
                    pass  # unsupported

                def ports():
                    for port in value:
                        cmd.extend(['--publish', port, '\\\n'])

                def networks():
                    for network in value:
                        cmd.extend(['--network', network if self.is_external_network(network) else self.project_prefix(network), '\\\n'])

                def volumes():
                    for volume in value:
                        src, dst = volume.split(':')

                        if src.startswith('.'):
                            src = src.replace('.', self.compose_base_dir, 1)

                        if src.startswith('/'):
                            cmd.extend(['--mount', 'type=bind,src={},dst={}'.format(src, dst), '\\\n'])
                        else:
                            cmd.extend(['--mount', 'src={},dst={}'.format(self.project_prefix(src), dst), '\\\n'])

                def environment():
                    if isinstance(value, dict):
                        for k, v in value.items():
                            cmd.extend(['--env', '{}={}'.format(k, v), '\\\n'])
                    else:
                        for env in value:
                            if env.startswith('constraint') or env.startswith('affinity'):
                                constraint = env.split(':', 2)[1]
                                cmd.extend(['--constraint', "'{}'".format(constraint), '\\\n'])
                            else:
                                cmd.extend(['--env', env, '\\\n'])

                def replicas():
                    cmd.extend(['--replicas', value, '\\\n'])

                def unsupported():
                    print('WARNING: unsupported parameter {}'.format(parameter))

                locals().get(parameter, unsupported)()

            cmd.extend(service_image)
            cmd.extend(service_command)

            self.call(' '.join(cmd))

        if services_to_start:
            self.start(services_to_start)

    def pull(self):
        nodes = self.call("docker node ls | grep Ready | awk -F'[[:space:]][[:space:]]+' '{print $2}'").rstrip().split('\n')

        threads = []

        for node in nodes:
            cmd = '; '.join(['docker -H tcp://{}:2375 pull {}'.format(node, self.services[service]['image']) for service in self.filtered_services])
            threads.append((node, threading.Thread(target=self.call, args=(cmd,))))

        for node, thread in threads:
            print('Pulling on node {}'.format(node))
            thread.start()

        for node, thread in threads:
            thread.join()
            print('Node {} - DONE'.format(node))

    def stop(self):
        services = filter(self.is_service_exists, self.filtered_services)
        cmd_args = ['{}={}'.format(self.project_prefix(service), 0) for service in services]
        if cmd_args:
            self.call('docker service scale ' + ' '.join(cmd_args))

    def rm(self):
        services = filter(self.is_service_exists, self.filtered_services)
        cmd_args = [self.project_prefix(service) for service in services]
        if cmd_args:
            self.call('docker service rm ' + ' '.join(cmd_args))

    def start(self, services=None):
        if services is None:
            services = self.filtered_services

        cmd = 'docker service scale ' + \
              ' '.join(['{}={}'.format(self.project_prefix(service), self.services[service].get('replicas', '1')) for service in services])
        self.call(cmd)

def main():
    parser = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=50, width=120))
    parser.add_argument('-f', '--file', type=argparse.FileType(), help='Specify an alternate compose file (default: docker-compose.yml)', default='docker-compose.yml')
    parser.add_argument('-p', '--project-name', help='Specify an alternate project name (default: directory name)')
    parser.add_argument('--dry-run', action='store_true')
    subparsers = parser.add_subparsers(title='Command')
    parser.add_argument('_service', metavar='service', nargs='*', help='List of services to run the command for')

    services_parser = argparse.ArgumentParser(add_help=False)
    services_parser.add_argument('service', nargs='*', help='List of services to run the command for')

    pull_parser = subparsers.add_parser('pull', help='Pull service images', add_help=False, parents=[services_parser])
    pull_parser.set_defaults(command='pull')

    rm_parser = subparsers.add_parser('rm', help='Stop and remove services', add_help=False, parents=[services_parser])
    rm_parser.set_defaults(command='rm')
    rm_parser.add_argument('-f', help='docker-compose compatibility; ignored', action='store_true')

    start_parser = subparsers.add_parser('start', help='Start services', add_help=False, parents=[services_parser])
    start_parser.set_defaults(command='start')

    stop_parser = subparsers.add_parser('stop', help='Stop services', add_help=False, parents=[services_parser])
    stop_parser.set_defaults(command='stop')

    up_parser = subparsers.add_parser('up', help='Create and start services', add_help=False, parents=[services_parser])
    up_parser.set_defaults(command='up')
    up_parser.add_argument('-d', help='docker-compose compatibility; ignored', action='store_true')

    args = parser.parse_args(sys.argv[1:])

    global debug
    debug = args.dry_run

    compose_base_dir = os.path.dirname(os.path.abspath(args.file.name))

    if not args.project_name:
        args.project_name = os.path.basename(compose_base_dir)

    docker_compose = DockerCompose(yaml.load(args.file, yodl.OrderedDictYAMLLoader), args.project_name, compose_base_dir + '/', args.service)
    getattr(docker_compose, args.command)()

if __name__ == "__main__":
    main()
<commit_msg>Quote env values<commit_after>#!/usr/bin/env python

import argparse
import os
import subprocess
import sys
import threading
from collections import OrderedDict

import yaml
import yodl

debug = False

class DockerCompose:
    def __init__(self, compose, project, compose_base_dir, requested_services):
        self.project = project
        self.compose_base_dir = compose_base_dir
        self.services = self.merge_services(compose.get('services', {}))
        self.networks = compose.get('networks', {})
        self.volumes = compose.get('volumes', {})
        self.filtered_services = filter(lambda service: not requested_services or service in requested_services, self.services)

    def project_prefix(self, value):
        return '{}_{}'.format(self.project, value)

    def merge_services(self, services):
        result = OrderedDict()

        for service in services:
            service_config = services[service]
            result[service] = service_config

            if 'extends' in service_config:
                extended_config = service_config['extends']
                extended_service = extended_config['service']

                del result[service]['extends']

                if 'file' in extended_config:
                    extended_service_data = self.merge_services(
                        yaml.load(open(self.compose_base_dir + extended_config['file'], 'r'), yodl.OrderedDictYAMLLoader)['services']
                    )[extended_service]
                else:
                    extended_service_data = result[extended_service]

                for k, v in extended_service_data.items():
                    if k not in result[service]:
                        result[service][k] = v

        return result

    @staticmethod
    def call(cmd, ignore_return_code=False):
        print('Running: \n' + cmd + '\n')
        if not debug:
            ps = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
            returncode = ps.wait()
            stdout = ps.communicate()[0]
            if returncode != 0 and not ignore_return_code:
                print >> sys.stderr, ('Error: command "{}" failed: {}'.format(cmd, stdout))
                sys.exit(returncode)
            else:
                return stdout

    def is_service_exists(self, service):
        return self.call('docker service ls -f name={} | tail -n +2'.format(self.project_prefix(service)))

    def is_external_network(self, network):
        return isinstance(self.networks[network], dict) and 'external' in self.networks[network]

    def up(self):
        for network in self.networks:
            if not self.is_external_network(network):
                cmd = '[ "`docker network ls -f name={0} | tail -n +2`" != "" ] || docker network create --driver overlay --opt encrypted {0}' \
                    .format(self.project_prefix(network))
                self.call(cmd)

        for volume in self.volumes:
            cmd = '[ "`docker volume ls -f name={0} | tail -n +2`" != "" ] || docker volume create --name {0}'.format(self.project_prefix(volume))
            self.call(cmd)

        services_to_start = []

        for service in self.filtered_services:
            if self.is_service_exists(service):
                services_to_start.append(service)
                continue

            service_config = self.services[service]
            cmd = ['docker service create --with-registry-auth \\\n --name', self.project_prefix(service), '\\\n']

            service_image = []
            service_command = []

            for parameter in service_config:
                value = service_config[parameter]

                def restart():
                    cmd.extend(['--restart-condition', {'always': 'any'}[value], '\\\n'])

                def logging():
                    cmd.extend(['--log-driver', value['driver'], '\\\n'])
                    log_opts = value['options']
                    for k, v in log_opts.items():
                        cmd.extend(['--log-opt', '{}={}'.format(k, v), '\\\n'])

                def mem_limit():
                    cmd.extend(['--limit-memory', value, '\\\n'])

                def image():
                    service_image.append(value)

                def command():
                    service_command.extend(value.split(' '))

                def expose():
                    pass  # unsupported

                def container_name():
                    pass  # unsupported

                def hostname():
                    pass  # unsupported; waiting for https://github.com/docker/docker/issues/24877

                def extra_hosts():
                    pass  # unsupported

                def ports():
                    for port in value:
                        cmd.extend(['--publish', port, '\\\n'])

                def networks():
                    for network in value:
                        cmd.extend(['--network', network if self.is_external_network(network) else self.project_prefix(network), '\\\n'])

                def volumes():
                    for volume in value:
                        src, dst = volume.split(':')

                        if src.startswith('.'):
                            src = src.replace('.', self.compose_base_dir, 1)

                        if src.startswith('/'):
                            cmd.extend(['--mount', 'type=bind,src={},dst={}'.format(src, dst), '\\\n'])
                        else:
                            cmd.extend(['--mount', 'src={},dst={}'.format(self.project_prefix(src), dst), '\\\n'])

                def environment():
                    if isinstance(value, dict):
                        for k, v in value.items():
                            cmd.extend(['--env', '{}="{}"'.format(k, v), '\\\n'])
                    else:
                        for env in value:
                            if env.startswith('constraint') or env.startswith('affinity'):
                                constraint = env.split(':', 2)[1]
                                cmd.extend(['--constraint', "'{}'".format(constraint), '\\\n'])
                            else:
                                cmd.extend(['--env', env, '\\\n'])

                def replicas():
                    cmd.extend(['--replicas', value, '\\\n'])

                def unsupported():
                    print('WARNING: unsupported parameter {}'.format(parameter))

                locals().get(parameter, unsupported)()

            cmd.extend(service_image)
            cmd.extend(service_command)

            self.call(' '.join(cmd))

        if services_to_start:
            self.start(services_to_start)

    def pull(self):
        nodes = self.call("docker node ls | grep Ready | awk -F'[[:space:]][[:space:]]+' '{print $2}'").rstrip().split('\n')

        threads = []

        for node in nodes:
            cmd = '; '.join(['docker -H tcp://{}:2375 pull {}'.format(node, self.services[service]['image']) for service in self.filtered_services])
            threads.append((node, threading.Thread(target=self.call, args=(cmd,))))

        for node, thread in threads:
            print('Pulling on node {}'.format(node))
            thread.start()

        for node, thread in threads:
            thread.join()
            print('Node {} - DONE'.format(node))

    def stop(self):
        services = filter(self.is_service_exists, self.filtered_services)
        cmd_args = ['{}={}'.format(self.project_prefix(service), 0) for service in services]
        if cmd_args:
            self.call('docker service scale ' + ' '.join(cmd_args))

    def rm(self):
        services = filter(self.is_service_exists, self.filtered_services)
        cmd_args = [self.project_prefix(service) for service in services]
        if cmd_args:
            self.call('docker service rm ' + ' '.join(cmd_args))

    def start(self, services=None):
        if services is None:
            services = self.filtered_services

        cmd = 'docker service scale ' + \
              ' '.join(['{}={}'.format(self.project_prefix(service), self.services[service].get('replicas', '1')) for service in services])
        self.call(cmd)

def main():
    parser = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=50, width=120))
    parser.add_argument('-f', '--file', type=argparse.FileType(), help='Specify an alternate compose file (default: docker-compose.yml)', default='docker-compose.yml')
    parser.add_argument('-p', '--project-name', help='Specify an alternate project name (default: directory name)')
    parser.add_argument('--dry-run', action='store_true')
    subparsers = parser.add_subparsers(title='Command')
    parser.add_argument('_service', metavar='service', nargs='*', help='List of services to run the command for')

    services_parser = argparse.ArgumentParser(add_help=False)
    services_parser.add_argument('service', nargs='*', help='List of services to run the command for')

    pull_parser = subparsers.add_parser('pull', help='Pull service images', add_help=False, parents=[services_parser])
    pull_parser.set_defaults(command='pull')

    rm_parser = subparsers.add_parser('rm', help='Stop and remove services', add_help=False, parents=[services_parser])
    rm_parser.set_defaults(command='rm')
    rm_parser.add_argument('-f', help='docker-compose compatibility; ignored', action='store_true')

    start_parser = subparsers.add_parser('start', help='Start services', add_help=False, parents=[services_parser])
    start_parser.set_defaults(command='start')

    stop_parser = subparsers.add_parser('stop', help='Stop services', add_help=False, parents=[services_parser])
    stop_parser.set_defaults(command='stop')

    up_parser = subparsers.add_parser('up', help='Create and start services', add_help=False, parents=[services_parser])
    up_parser.set_defaults(command='up')
    up_parser.add_argument('-d', help='docker-compose compatibility; ignored', action='store_true')

    args = parser.parse_args(sys.argv[1:])

    global debug
    debug = args.dry_run

    compose_base_dir = os.path.dirname(os.path.abspath(args.file.name))

    if not args.project_name:
        args.project_name = os.path.basename(compose_base_dir)

    docker_compose = DockerCompose(yaml.load(args.file, yodl.OrderedDictYAMLLoader), args.project_name, compose_base_dir + '/', args.service)
    getattr(docker_compose, args.command)()

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>import atexit
import logging
import psutil
import shlex
import signal
import subprocess
import time

from .timeout import timeout, TimeoutError

log = logging.getLogger(__name__)

COMMANDS = {
    'unicorn': 'unicorn -D -P "{pidfile}" {args}',
    'unicorn_rails': 'unicorn_rails -D {args}',
    'unicorn_bin': '{unicorn_bin} -D -P "{pidfile}" {args}',
    'gunicorn': 'gunicorn -D -p "{pidfile}" {args}',
    'gunicorn_django': 'gunicorn_django -D -p "{pidfile}" {args}',
    'gunicorn_bin': '{gunicorn_bin} -D -p "{pidfile}" {args}'
}

MANAGED_PIDS = set([])

class HerderError(Exception):
    pass

class Herder(object):
    """

    The Herder class manages a single unicorn instance and its worker
    children. It has few configuration options: you simply instantiate a
    Herder, spawn unicorn with ``spawn()``, and then start a monitoring loop
    with the ``loop()`` method.

    The ``loop()`` method will exit with a status code, which by default will
    be used as the exit status of the ``unicornherder`` command line utility.

    Example::

        herder = Herder()
        if herder.spawn():
            sys.exit(herder.loop())

    """

    def __init__(self, unicorn='gunicorn', unicorn_bin=None, gunicorn_bin=None,
                 pidfile=None, boot_timeout=30, overlap=120, args=''):
        """

        Creates a new Herder instance.

        unicorn      - the type of unicorn to herd; either 'unicorn' or 'gunicorn'
                       (Default: gunicorn)
        unicorn_bin  - path of specific unicorn to run
                       (Default: None)
        gunicorn_bin - path of specific gunicorn to run
                       (Default: None)
        pidfile      - path of the pidfile to write
                       (Default: gunicorn.pid or unicorn.pid depending on the value of
                        the unicorn parameter)
        boot_timeout - how long to wait for the new process to daemonize itself
        overlap      - how long to wait before killing the old unicorns when reloading
        args         - any additional arguments to pass to the unicorn executable
                       (Default: '')

        """

        self.unicorn_bin = unicorn_bin
        self.gunicorn_bin = gunicorn_bin

        if unicorn_bin:
            self.unicorn = unicorn_bin
        elif gunicorn_bin:
            self.unicorn = gunicorn_bin
        else:
            self.unicorn = unicorn
        self.pidfile = '%s.pid' % self.unicorn if pidfile is None else pidfile
        self.args = args
        self.boot_timeout = boot_timeout
        self.overlap = overlap

        try:
            if not unicorn_bin and not gunicorn_bin:
                COMMANDS[self.unicorn]
        except KeyError:
            raise HerderError('Unknown unicorn type: %s' % self.unicorn)

        self.master = None
        self.reloading = False
        self.terminating = False

    def spawn(self):
        """

        Spawn a new unicorn instance.

        Returns False if unicorn fails to daemonize, and True otherwise.

        """
        if self.unicorn in COMMANDS:
            cmd = COMMANDS[self.unicorn]
            cmd = cmd.format(pidfile=self.pidfile, args=self.args)
        elif self.unicorn_bin:
            cmd = COMMANDS['unicorn_bin']
            cmd = cmd.format(unicorn_bin=self.unicorn, pidfile=self.pidfile, args=self.args)
        elif self.gunicorn_bin:
            cmd = COMMANDS['gunicorn_bin']
            cmd = cmd.format(gunicorn_bin=self.unicorn, pidfile=self.pidfile, args=self.args)
        else:
            return False

        log.debug("Calling %s: %s", self.unicorn, cmd)

        cmd = shlex.split(cmd)
        try:
            process = subprocess.Popen(cmd)
        except OSError as e:
            if e.errno == 2:
                log.error("Command '%s' not found. Is it installed?", cmd[0])
                return False
            else:
                raise

        MANAGED_PIDS.add(process.pid)

        try:
            with timeout(self.boot_timeout):
                process.wait()
        except TimeoutError:
            log.error('%s failed to daemonize within %s seconds. Sending TERM '
                      'and exiting.', self.unicorn, self.boot_timeout)
            if process.poll() is None:
                process.terminate()
            return False

        # If we got this far, unicorn has daemonized, and we no longer need to
        # worry about the original process.
        MANAGED_PIDS.remove(process.pid)

        # The unicorn herder does a graceful unicorn restart on HUP
        signal.signal(signal.SIGHUP, self._handle_HUP)

        # Forward other useful signals to the currently tracked master
        # process.
        #
        # We do NOT forward SIGWINCH, because it is triggered by terminal
        # resize, leading to some *seriously* weird behaviour (resize
        # xterm, unicorn workers are killed).
        for sig in ['INT', 'QUIT', 'TERM', 'TTIN', 'TTOU', 'USR1', 'USR2']:
            signal.signal(getattr(signal, 'SIG%s' % sig),
                          self._handle_signal(sig))

        return True

    def loop(self):
        """Enter the monitoring loop"""
        while True:
            if not self._loop_inner():
                # The unicorn has died. So should we.
                log.error('%s died. Exiting.', self.unicorn.title())
                return 1
            time.sleep(2)

    def _loop_inner(self):
        old_master = self.master
        pid = self._read_pidfile()

        if pid is None:
            return False

        try:
            self.master = psutil.Process(pid)
        except psutil.NoSuchProcess:
            return False

        if old_master is None:
            log.info('%s booted (PID %s)', self.unicorn.title(), self.master.pid)

            MANAGED_PIDS.add(self.master.pid)

        # Unicorn has forked a new master
        if old_master is not None and self.master.pid != old_master.pid:
            log.info('%s changed PID (was %s, now %s)',
                     self.unicorn.title(),
                     old_master.pid,
                     self.master.pid)

            MANAGED_PIDS.add(self.master.pid)

            if self.reloading:
                _wait_for_workers(self.overlap)
                _kill_old_master(old_master)
                self.reloading = False

            MANAGED_PIDS.remove(old_master.pid)

        return True

    def _read_pidfile(self):
        for _ in range(5):
            try:
                content = open(self.pidfile).read()
            except IOError as e:
                try:
                    log.debug('pidfile missing, checking for %s.oldbin', self.pidfile)
                    content = open(self.pidfile + ".oldbin").read()
                except IOError as e:
                    # If we are expecting unicorn to die, then this is normal, and
                    # we can just return None, thus triggering a clean exit of the
                    # Herder.
                    if self.terminating:
                        return None
                    else:
                        log.debug('Got IOError while attempting to read pidfile: %s', e)
                        log.debug('This is usually not fatal. Retrying in a moment...')
                        time.sleep(1)
                        continue
            try:
                pid = int(content)
            except ValueError as e:
                log.debug('Got ValueError while reading pidfile. Is "%s" an integer? %s',
                          content, e)
                log.debug('This is usually not fatal. Retrying in a moment...')
                time.sleep(1)
                continue

            return pid

        raise HerderError('Failed to read pidfile %s after 5 attempts, aborting!' %
                          self.pidfile)

    def _handle_signal(self, name):
        def _handler(signum, frame):
            if self.master is None:
                log.warn("Caught %s but have no tracked process.", name)
                return

            if signum in [signal.SIGINT, signal.SIGQUIT, signal.SIGTERM]:
                log.debug("Caught %s: expecting termination.", name)
                self.terminating = True

            log.debug("Forwarding %s to PID %s", name, self.master.pid)
            self.master.send_signal(signum)

        return _handler

    def _handle_HUP(self, signum, frame):
        if self.master is None:
            log.warn("Caught HUP but have no tracked process.")
            return

        log.info("Caught HUP: gracefully restarting PID %s", self.master.pid)
        self.reloading = True
        self.master.send_signal(signal.SIGUSR2)

#
# If the unicorn herder exits abnormally, it is essential that unicorn
# dies as well. Register an atexit callback to kill off any surviving
# unicorns.
#
@atexit.register
def _emergency_slaughter():
    for pid in MANAGED_PIDS:
        try:
            proc = psutil.Process(pid)
            proc.kill()
        except:
            pass

def _wait_for_workers(overlap):
    # TODO: do something smarter here
    time.sleep(overlap)

def _kill_old_master(process):
    # SIGWINCH tells the master to kill its workers:
    log.debug("Sending WINCH to old master (PID %s)", process.pid)
    process.send_signal(signal.SIGWINCH)
    time.sleep(1)
    log.debug("Sending QUIT to old master (PID %s)", process.pid)
    process.send_signal(signal.SIGQUIT)
<commit_msg>Document the behaviour of _kill_old_master<commit_after>import atexit
import logging
import psutil
import shlex
import signal
import subprocess
import time

from .timeout import timeout, TimeoutError

log = logging.getLogger(__name__)

COMMANDS = {
    'unicorn': 'unicorn -D -P "{pidfile}" {args}',
    'unicorn_rails': 'unicorn_rails -D {args}',
    'unicorn_bin': '{unicorn_bin} -D -P "{pidfile}" {args}',
    'gunicorn': 'gunicorn -D -p "{pidfile}" {args}',
    'gunicorn_django': 'gunicorn_django -D -p "{pidfile}" {args}',
    'gunicorn_bin': '{gunicorn_bin} -D -p "{pidfile}" {args}'
}

MANAGED_PIDS = set([])

class HerderError(Exception):
    pass

class Herder(object):
    """

    The Herder class manages a single unicorn instance and its worker
    children. It has few configuration options: you simply instantiate a
    Herder, spawn unicorn with ``spawn()``, and then start a monitoring loop
    with the ``loop()`` method.

    The ``loop()`` method will exit with a status code, which by default will
    be used as the exit status of the ``unicornherder`` command line utility.

    Example::

        herder = Herder()
        if herder.spawn():
            sys.exit(herder.loop())

    """

    def __init__(self, unicorn='gunicorn', unicorn_bin=None, gunicorn_bin=None,
                 pidfile=None, boot_timeout=30, overlap=120, args=''):
        """

        Creates a new Herder instance.

        unicorn      - the type of unicorn to herd; either 'unicorn' or 'gunicorn'
                       (Default: gunicorn)
        unicorn_bin  - path of specific unicorn to run
                       (Default: None)
        gunicorn_bin - path of specific gunicorn to run
                       (Default: None)
        pidfile      - path of the pidfile to write
                       (Default: gunicorn.pid or unicorn.pid depending on the value of
                        the unicorn parameter)
        boot_timeout - how long to wait for the new process to daemonize itself
        overlap      - how long to wait before killing the old unicorns when reloading
        args         - any additional arguments to pass to the unicorn executable
                       (Default: '')

        """

        self.unicorn_bin = unicorn_bin
        self.gunicorn_bin = gunicorn_bin

        if unicorn_bin:
            self.unicorn = unicorn_bin
        elif gunicorn_bin:
            self.unicorn = gunicorn_bin
        else:
            self.unicorn = unicorn
        self.pidfile = '%s.pid' % self.unicorn if pidfile is None else pidfile
        self.args = args
        self.boot_timeout = boot_timeout
        self.overlap = overlap

        try:
            if not unicorn_bin and not gunicorn_bin:
                COMMANDS[self.unicorn]
        except KeyError:
            raise HerderError('Unknown unicorn type: %s' % self.unicorn)

        self.master = None
        self.reloading = False
        self.terminating = False

    def spawn(self):
        """

        Spawn a new unicorn instance.

        Returns False if unicorn fails to daemonize, and True otherwise.

        """
        if self.unicorn in COMMANDS:
            cmd = COMMANDS[self.unicorn]
            cmd = cmd.format(pidfile=self.pidfile, args=self.args)
        elif self.unicorn_bin:
            cmd = COMMANDS['unicorn_bin']
            cmd = cmd.format(unicorn_bin=self.unicorn, pidfile=self.pidfile, args=self.args)
        elif self.gunicorn_bin:
            cmd = COMMANDS['gunicorn_bin']
            cmd = cmd.format(gunicorn_bin=self.unicorn, pidfile=self.pidfile, args=self.args)
        else:
            return False

        log.debug("Calling %s: %s", self.unicorn, cmd)

        cmd = shlex.split(cmd)
        try:
            process = subprocess.Popen(cmd)
        except OSError as e:
            if e.errno == 2:
                log.error("Command '%s' not found. Is it installed?", cmd[0])
                return False
            else:
                raise

        MANAGED_PIDS.add(process.pid)

        try:
            with timeout(self.boot_timeout):
                process.wait()
        except TimeoutError:
            log.error('%s failed to daemonize within %s seconds. Sending TERM '
                      'and exiting.', self.unicorn, self.boot_timeout)
            if process.poll() is None:
                process.terminate()
            return False

        # If we got this far, unicorn has daemonized, and we no longer need to
        # worry about the original process.
        MANAGED_PIDS.remove(process.pid)

        # The unicorn herder does a graceful unicorn restart on HUP
        signal.signal(signal.SIGHUP, self._handle_HUP)

        # Forward other useful signals to the currently tracked master
        # process.
        #
        # We do NOT forward SIGWINCH, because it is triggered by terminal
        # resize, leading to some *seriously* weird behaviour (resize
        # xterm, unicorn workers are killed).
        for sig in ['INT', 'QUIT', 'TERM', 'TTIN', 'TTOU', 'USR1', 'USR2']:
            signal.signal(getattr(signal, 'SIG%s' % sig),
                          self._handle_signal(sig))

        return True

    def loop(self):
        """Enter the monitoring loop"""
        while True:
            if not self._loop_inner():
                # The unicorn has died. So should we.
                log.error('%s died. Exiting.', self.unicorn.title())
                return 1
            time.sleep(2)

    def _loop_inner(self):
        old_master = self.master
        pid = self._read_pidfile()

        if pid is None:
            return False

        try:
            self.master = psutil.Process(pid)
        except psutil.NoSuchProcess:
            return False

        if old_master is None:
            log.info('%s booted (PID %s)', self.unicorn.title(), self.master.pid)

            MANAGED_PIDS.add(self.master.pid)

        # Unicorn has forked a new master
        if old_master is not None and self.master.pid != old_master.pid:
            log.info('%s changed PID (was %s, now %s)',
                     self.unicorn.title(),
                     old_master.pid,
                     self.master.pid)

            MANAGED_PIDS.add(self.master.pid)

            if self.reloading:
                _wait_for_workers(self.overlap)
                _kill_old_master(old_master)
                self.reloading = False

            MANAGED_PIDS.remove(old_master.pid)

        return True

    def _read_pidfile(self):
        for _ in range(5):
            try:
                content = open(self.pidfile).read()
            except IOError as e:
                try:
                    log.debug('pidfile missing, checking for %s.oldbin', self.pidfile)
                    content = open(self.pidfile + ".oldbin").read()
                except IOError as e:
                    # If we are expecting unicorn to die, then this is normal, and
                    # we can just return None, thus triggering a clean exit of the
                    # Herder.
                    if self.terminating:
                        return None
                    else:
                        log.debug('Got IOError while attempting to read pidfile: %s', e)
                        log.debug('This is usually not fatal. Retrying in a moment...')
                        time.sleep(1)
                        continue
            try:
                pid = int(content)
            except ValueError as e:
                log.debug('Got ValueError while reading pidfile. Is "%s" an integer? %s',
                          content, e)
                log.debug('This is usually not fatal. Retrying in a moment...')
                time.sleep(1)
                continue

            return pid

        raise HerderError('Failed to read pidfile %s after 5 attempts, aborting!' %
                          self.pidfile)

    def _handle_signal(self, name):
        def _handler(signum, frame):
            if self.master is None:
                log.warn("Caught %s but have no tracked process.", name)
                return

            if signum in [signal.SIGINT, signal.SIGQUIT, signal.SIGTERM]:
                log.debug("Caught %s: expecting termination.", name)
                self.terminating = True

            log.debug("Forwarding %s to PID %s", name, self.master.pid)
            self.master.send_signal(signum)

        return _handler

    def _handle_HUP(self, signum, frame):
        if self.master is None:
            log.warn("Caught HUP but have no tracked process.")
            return

        log.info("Caught HUP: gracefully restarting PID %s", self.master.pid)
        self.reloading = True
        self.master.send_signal(signal.SIGUSR2)

#
# If the unicorn herder exits abnormally, it is essential that unicorn
# dies as well. Register an atexit callback to kill off any surviving
# unicorns.
#
@atexit.register
def _emergency_slaughter():
    for pid in MANAGED_PIDS:
        try:
            proc = psutil.Process(pid)
            proc.kill()
        except:
            pass

def _wait_for_workers(overlap):
    # TODO: do something smarter here
    time.sleep(overlap)

def _kill_old_master(process):
    """Shut down the old server gracefully.

    There's a bit of extra complexity here, because Unicorn and Gunicorn handle
    signals differently : both respond to SIGWINCH by gracefully stopping their
    workers, but while Unicorn treats SIGQUIT as a graceful shutdown and
    SIGTERM as a quick shutdown, Gunicorn reverses the meaning of these two.

    <http://unicorn.bogomips.org/SIGNALS.html>
    <http://gunicorn-docs.readthedocs.org/en/latest/signals.html>

    We get around this by sending SIGWINCH first, giving the worker processes
    some time to shut themselves down first.

    """
    log.debug("Sending WINCH to old master (PID %s)", process.pid)
    process.send_signal(signal.SIGWINCH)
    time.sleep(1)
    log.debug("Sending QUIT to old master (PID %s)", process.pid)
    process.send_signal(signal.SIGQUIT)
<|endoftext|>


<commit_before>import subprocess

import rdflib

from django.http import Http404, HttpResponse, HttpResponseBadRequest
from django.template import RequestContext, loader
from django.template.defaultfilters import slugify
from django.shortcuts import render_to_response

from humfrey.linkeddata.views import RDFView
from humfrey.utils.resource import Resource
from humfrey.utils.views import renderer
from humfrey.utils.namespaces import expand, NS

class GraphVizView(RDFView):
    _QUERY = """
      CONSTRUCT {
        %(page_uri)s foaf:topic ?entity .
        ?entity a ?type %(propertyTriples)s .
          %(relationPattern)s .
      } WHERE {
        { SELECT DISTINCT ?entity WHERE { %(subject)s (%(relationAlternation)s){0,%(depth)d} %(object)s } } .
        ?entity a ?type .
        OPTIONAL { %(relationPattern)s } .
        %(propertyPatterns)s
        OPTIONAL { ?entity rdfs:label ?title } .
        FILTER (?relation in (%(relationList)s)) .
        %(excludeTypesFilter)s
      }
    """

    def handle_GET(self, request, context, root=None, relations=None, template='graphviz/graphviz', depth=4, max_depth=5, exclude_types=None, properties=None, inverted=None):
        root = expand(root or request.GET.get('root', ''))
        relations = relations or [expand(relation) for relation in request.GET.getlist('relation')]
        exclude_types = exclude_types or [expand(t) for t in request.GET.getlist('exclude_type')]
        properties = properties or [expand(p) for p in request.GET.getlist('property')]
        try:
            depth = min(int(request.GET.get('depth', depth)), max_depth)
        except (TypeError, ValueError):
            return HttpResponseBadRequest()

        inverted = inverted if (inverted is not None) else request.GET.get('inverted') == 'true'
        if inverted:
            subj, obj = '?entity', root.n3()
            relation_pattern = '?entity ?relation ?parent'
        else:
            subj, obj = root.n3(), '?entity'
            relation_pattern = '?parent ?relation ?entity'

        types = self.get_types(root)
        if not types or not relations:
            raise Http404

        page_uri = rdflib.URIRef(request.build_absolute_uri())

        query = self._QUERY % {'subject': subj,
                               'object': obj,
                               'depth': depth,
                               'relationList': ', '.join(r.n3() for r in relations),
                               'excludeTypesFilter': ('FILTER (?type not in (%s)) .' % ', '.join(t.n3() for t in exclude_types)) if exclude_types else '',
                               'relationAlternation': '|'.join(r.n3() for r in relations),
                               'relationPattern': relation_pattern,
                               'page_uri': page_uri.n3(),
                               'propertyPatterns': '\n        '.join('OPTIONAL { ?entity %s ?p%s } .' % (p.n3(), i) for i, p in enumerate(properties)),
                               'propertyTriples': ''.join(';\n          %s ?p%s' % (p.n3(), i) for i, p in enumerate(properties))
                               }
        print query
        graph = self.endpoint.query(query)

        context.update({
            'graph': graph,
            'queries': [graph.query],
            'subjects': [Resource(s, graph, self.endpoint) for s in set(graph.objects(page_uri, NS['foaf'].topic))],
            'subject': Resource(root, graph, self.endpoint),
            'inverted': inverted,
            'relations': relations,
        })

        for subject in context['subjects']:

            if not inverted:
                subject.children = set(Resource(s, graph, self.endpoint) for relation in relations for s in graph.objects(subject._identifier, relation))
            else:
                subject.children = set(Resource(s, graph, self.endpoint) for relation in relations for s in graph.subjects(relation, subject._identifier))
            for child in subject.children:
                if (page_uri, NS['foaf'].topic, child._identifier) in graph:
                    child.display = True

        return self.render(request, context, template)

    _DOT_LAYOUTS = "circo dot fdp neato nop nop1 nop2 osage patchwork sfdp twopi".split()
    _DOT_OUTPUTS = [
        dict(format='bmp', mimetypes=('image/x-bmp','image/x-ms-bmp'), name='BMP', dot_output='bmp'),
        dict(format='xdot', mimetypes=('text/vnd.graphviz',), name='xDOT', dot_output='xdot', priority=0.9),
        dict(format='gv', mimetypes=('text/vnd.graphviz',), name='DOT (GraphViz)', dot_output='gv'),
        dict(format='jpeg', mimetypes=('image/jpeg',), name='JPEG', dot_output='jpeg'),
        dict(format='png', mimetypes=('image/png',), name='PNG', dot_output='png'),
        dict(format='ps', mimetypes=('application/postscript',), name='PostScript', dot_output='ps'),
        dict(format='pdf', mimetypes=('application/pdf',), name='PDF', dot_output='pdf'),
        dict(format='svg', mimetypes=('image/svg+xml',), name='SVG', dot_output='svg'),
    ]

    def _get_dot_renderer(output):
        def dot_renderer(self, request, context, template_name):
            layout = request.GET.get('layout')
            if layout not in self._DOT_LAYOUTS:
                layout = 'fdp'
            template = loader.get_template(template_name + '.gv')
            plain_gv = template.render(RequestContext(request, context))
            dot = subprocess.Popen(['dot', '-K'+layout, '-T'+dot_output], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
            dot_stdout, _ = dot.communicate(input=plain_gv.encode('utf-8'))
            response = HttpResponse(dot_stdout, mimetype=output['mimetypes'][0])
            response['Content-Disposition'] = 'inline; filename="%s.%s"' % (slugify(context['subject'].dcterms_title)[:32], output['format'])
            return response

        dot_output = output.pop('dot_output')
        dot_renderer.__name__ = 'render_%s' % output['format']
        return renderer(**output)(dot_renderer)

    for output in _DOT_OUTPUTS:
        locals()['render_%s' % output['format']] = _get_dot_renderer(output)
    del _get_dot_renderer, output

    @renderer(format="gv", mimetypes=('text/vnd.graphviz',), name="DOT (GraphViz)")
    def render_gv(self, request, context, template_name):
        layout = request.GET.get('layout')
        if layout not in self._DOT_LAYOUTS:
            layout = 'fdp'
        template = loader.get_template(template_name + '.gv')
        plain_gv = template.render(RequestContext(request, context))
        dot = subprocess.Popen(['dot', '-K'+layout, '-Txdot'], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
        dot_stdout, _ = dot.communicate(input=plain_gv.encode('utf-8'))
        response = HttpResponse(dot_stdout, mimetype='text/vnd.graphviz')
        response['Content-Disposition'] = 'attachment; filename="%s.gv"' % slugify(context['subject'].dcterms_title)[:32]
        return response

    @renderer(format="graphml", mimetypes=('application/x-graphml+xml',), name="GraphML")
    def render_graphml(self, request, context, template_name):
        response = render_to_response(template_name + '.graphml',
                                      context, context_instance=RequestContext(request),
                                      mimetype='application/x-graphml+xml')
        response['Content-Disposition'] = 'attachment; filename="%s.graphml"' % slugify(context['subject'].dcterms_title)[:32]
        return response<commit_msg>Removed debugging print statement.<commit_after>import subprocess

import rdflib

from django.http import Http404, HttpResponse, HttpResponseBadRequest
from django.template import RequestContext, loader
from django.template.defaultfilters import slugify
from django.shortcuts import render_to_response

from humfrey.linkeddata.views import RDFView
from humfrey.utils.resource import Resource
from humfrey.utils.views import renderer
from humfrey.utils.namespaces import expand, NS

class GraphVizView(RDFView):
    _QUERY = """
      CONSTRUCT {
        %(page_uri)s foaf:topic ?entity .
        ?entity a ?type %(propertyTriples)s .
          %(relationPattern)s .
      } WHERE {
        { SELECT DISTINCT ?entity WHERE { %(subject)s (%(relationAlternation)s){0,%(depth)d} %(object)s } } .
        ?entity a ?type .
        OPTIONAL { %(relationPattern)s } .
        %(propertyPatterns)s
        OPTIONAL { ?entity rdfs:label ?title } .
        FILTER (?relation in (%(relationList)s)) .
        %(excludeTypesFilter)s
      }
    """

    def handle_GET(self, request, context, root=None, relations=None, template='graphviz/graphviz', depth=4, max_depth=5, exclude_types=None, properties=None, inverted=None):
        root = expand(root or request.GET.get('root', ''))
        relations = relations or [expand(relation) for relation in request.GET.getlist('relation')]
        exclude_types = exclude_types or [expand(t) for t in request.GET.getlist('exclude_type')]
        properties = properties or [expand(p) for p in request.GET.getlist('property')]
        try:
            depth = min(int(request.GET.get('depth', depth)), max_depth)
        except (TypeError, ValueError):
            return HttpResponseBadRequest()

        inverted = inverted if (inverted is not None) else request.GET.get('inverted') == 'true'
        if inverted:
            subj, obj = '?entity', root.n3()
            relation_pattern = '?entity ?relation ?parent'
        else:
            subj, obj = root.n3(), '?entity'
            relation_pattern = '?parent ?relation ?entity'

        types = self.get_types(root)
        if not types or not relations:
            raise Http404

        page_uri = rdflib.URIRef(request.build_absolute_uri())

        query = self._QUERY % {'subject': subj,
                               'object': obj,
                               'depth': depth,
                               'relationList': ', '.join(r.n3() for r in relations),
                               'excludeTypesFilter': ('FILTER (?type not in (%s)) .' % ', '.join(t.n3() for t in exclude_types)) if exclude_types else '',
                               'relationAlternation': '|'.join(r.n3() for r in relations),
                               'relationPattern': relation_pattern,
                               'page_uri': page_uri.n3(),
                               'propertyPatterns': '\n        '.join('OPTIONAL { ?entity %s ?p%s } .' % (p.n3(), i) for i, p in enumerate(properties)),
                               'propertyTriples': ''.join(';\n          %s ?p%s' % (p.n3(), i) for i, p in enumerate(properties))
                               }
        graph = self.endpoint.query(query)

        context.update({
            'graph': graph,
            'queries': [graph.query],
            'subjects': [Resource(s, graph, self.endpoint) for s in set(graph.objects(page_uri, NS['foaf'].topic))],
            'subject': Resource(root, graph, self.endpoint),
            'inverted': inverted,
            'relations': relations,
        })

        for subject in context['subjects']:

            if not inverted:
                subject.children = set(Resource(s, graph, self.endpoint) for relation in relations for s in graph.objects(subject._identifier, relation))
            else:
                subject.children = set(Resource(s, graph, self.endpoint) for relation in relations for s in graph.subjects(relation, subject._identifier))
            for child in subject.children:
                if (page_uri, NS['foaf'].topic, child._identifier) in graph:
                    child.display = True

        return self.render(request, context, template)

    _DOT_LAYOUTS = "circo dot fdp neato nop nop1 nop2 osage patchwork sfdp twopi".split()
    _DOT_OUTPUTS = [
        dict(format='bmp', mimetypes=('image/x-bmp','image/x-ms-bmp'), name='BMP', dot_output='bmp'),
        dict(format='xdot', mimetypes=('text/vnd.graphviz',), name='xDOT', dot_output='xdot', priority=0.9),
        dict(format='gv', mimetypes=('text/vnd.graphviz',), name='DOT (GraphViz)', dot_output='gv'),
        dict(format='jpeg', mimetypes=('image/jpeg',), name='JPEG', dot_output='jpeg'),
        dict(format='png', mimetypes=('image/png',), name='PNG', dot_output='png'),
        dict(format='ps', mimetypes=('application/postscript',), name='PostScript', dot_output='ps'),
        dict(format='pdf', mimetypes=('application/pdf',), name='PDF', dot_output='pdf'),
        dict(format='svg', mimetypes=('image/svg+xml',), name='SVG', dot_output='svg'),
    ]

    def _get_dot_renderer(output):
        def dot_renderer(self, request, context, template_name):
            layout = request.GET.get('layout')
            if layout not in self._DOT_LAYOUTS:
                layout = 'fdp'
            template = loader.get_template(template_name + '.gv')
            plain_gv = template.render(RequestContext(request, context))
            dot = subprocess.Popen(['dot', '-K'+layout, '-T'+dot_output], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
            dot_stdout, _ = dot.communicate(input=plain_gv.encode('utf-8'))
            response = HttpResponse(dot_stdout, mimetype=output['mimetypes'][0])
            response['Content-Disposition'] = 'inline; filename="%s.%s"' % (slugify(context['subject'].dcterms_title)[:32], output['format'])
            return response

        dot_output = output.pop('dot_output')
        dot_renderer.__name__ = 'render_%s' % output['format']
        return renderer(**output)(dot_renderer)

    for output in _DOT_OUTPUTS:
        locals()['render_%s' % output['format']] = _get_dot_renderer(output)
    del _get_dot_renderer, output

    @renderer(format="gv", mimetypes=('text/vnd.graphviz',), name="DOT (GraphViz)")
    def render_gv(self, request, context, template_name):
        layout = request.GET.get('layout')
        if layout not in self._DOT_LAYOUTS:
            layout = 'fdp'
        template = loader.get_template(template_name + '.gv')
        plain_gv = template.render(RequestContext(request, context))
        dot = subprocess.Popen(['dot', '-K'+layout, '-Txdot'], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
        dot_stdout, _ = dot.communicate(input=plain_gv.encode('utf-8'))
        response = HttpResponse(dot_stdout, mimetype='text/vnd.graphviz')
        response['Content-Disposition'] = 'attachment; filename="%s.gv"' % slugify(context['subject'].dcterms_title)[:32]
        return response

    @renderer(format="graphml", mimetypes=('application/x-graphml+xml',), name="GraphML")
    def render_graphml(self, request, context, template_name):
        response = render_to_response(template_name + '.graphml',
                                      context, context_instance=RequestContext(request),
                                      mimetype='application/x-graphml+xml')
        response['Content-Disposition'] = 'attachment; filename="%s.graphml"' % slugify(context['subject'].dcterms_title)[:32]
        return response<|endoftext|>


<commit_before>import math
from sys import exit
import numpy as np
import uuid
from .onshape_api.client import Client
from .config import config, configFile
from colorama import Fore, Back, Style

# OnShape API client
client = Client(logging=False, creds=configFile)
client.useCollisionsConfigurations = config['useCollisionsConfigurations']

# If a versionId is provided, it will be used, else the main workspace is retrieved
if config['versionId'] == '':
    print("\n" + Style.BRIGHT + '* Retrieving workspace ID ...' + Style.RESET_ALL)
    document = client.get_document(config['documentId']).json()
    workspaceId = document['defaultWorkspace']['id']
    print(Fore.GREEN + "+ Using workspace id: " + workspaceId + Style.RESET_ALL)
else:
    print("\n" + Style.BRIGHT + '* Using configuration version ID ' +
          config['versionId']+' ...' + Style.RESET_ALL)

# Now, finding the assembly, according to given name in configuration, or else the first possible one
print("\n" + Style.BRIGHT +
      '* Retrieving elements in the document, searching for the assembly...' + Style.RESET_ALL)
if config['versionId'] == '':
    elements = client.list_elements(config['documentId'], workspaceId).json()
else:
    elements = client.list_elements(
        config['documentId'], config['versionId'], 'v').json()
assemblyId = None
assemblyName = ''
for element in elements:
    if element['type'] == 'Assembly' and \
            (config['assemblyName'] is False or element['name'] == config['assemblyName']):
        print(Fore.GREEN + "+ Found assembly, id: " +
              element['id']+', name: "'+element['name']+'"' + Style.RESET_ALL)
        assemblyName = element['name']
        assemblyId = element['id']

if assemblyId == None:
    print(Fore.RED + "ERROR: Unable to find assembly in this document" + Style.RESET_ALL)
    exit(1)

# Retrieving the assembly
print("\n" + Style.BRIGHT + '* Retrieving assembly "' +
      assemblyName+'" with id '+assemblyId + Style.RESET_ALL)
if config['versionId'] == '':
    assembly = client.get_assembly(
        config['documentId'], workspaceId, assemblyId)
else:
    assembly = client.get_assembly(
        config['documentId'], config['versionId'], assemblyId, 'v')

root = assembly['rootAssembly']

# Finds a (leaf) instance given the full path, typically A B C where A and B would be subassemblies and C
# the final part

def findInstance(path, instances=None):
    global assembly

    if instances is None:
        instances = assembly['rootAssembly']['instances']

    for instance in instances:
        if instance['id'] == path[0]:
            if len(path) == 1:
                # If the length of remaining path is 1, the part is in the current assembly/subassembly
                return instance
            else:
                # Else, we need to find the matching sub assembly to find the proper part (recursively)
                d = instance['documentId']
                m = instance['documentMicroversion']
                e = instance['elementId']
                for asm in assembly['subAssemblies']:
                    if asm['documentId'] == d and asm['documentMicroversion'] == m and asm['elementId'] == e:
                        return findInstance(path[1:], asm['instances'])

    print(Fore.RED + 'Could not find instance for ' + str(path) + Style.RESET_ALL)

# Collecting occurrences, the path is the assembly / sub assembly chain
occurrences = {}
for occurrence in root['occurrences']:
    occurrence['assignation'] = None
    occurrence['instance'] = findInstance(occurrence['path'])
    occurrence['transform'] = np.matrix(
        np.reshape(occurrence['transform'], (4, 4)))
    occurrence['linkName'] = None
    occurrences[tuple(occurrence['path'])] = occurrence

# Gets an occurrence given its full path

def getOccurrence(path):
    return occurrences[tuple(path)]

# Assignations are pieces that will be in the same link. Note that this is only for top-level
# item of the path (all sub assemblies and parts in assemblies are naturally in the same link as
# the parent), but other parts that can be connected with mates in top assemblies are then assigned to
# the link
assignations = {}

# Frames (mated with frame_ name) will be special links in the output file allowing to track some specific
# manually identified frames
frames = {}

def assignParts(root, parent):
    assignations[root] = parent
    for occurrence in occurrences.values():
        if occurrence['path'][0] == root:
            occurrence['assignation'] = parent

from .features import init as features_init, getLimits
features_init(client, config, root, workspaceId, assemblyId)

# First, features are scanned to find the DOFs. Links that they connects are then tagged
print("\n" + Style.BRIGHT +
      '* Getting assembly features, scanning for DOFs...' + Style.RESET_ALL)
trunk = None
relations = {}
features = root['features']
for feature in features:
    if feature['featureType'] == 'mateConnector':
        name = feature['featureData']['name']
        if name[0:5] == 'link_':
            name = name[5:]
            occurrences[(feature['featureData']['occurrence'][0],)
                        ]['linkName'] = name
    else:
        if feature['suppressed']:
            continue

        data = feature['featureData']

        if len(data['matedEntities']) != 2 or \
                len(data['matedEntities'][0]['matedOccurrence']) == 0 \
                or len(data['matedEntities'][1]['matedOccurrence']) == 0:
            continue

        child = data['matedEntities'][0]['matedOccurrence'][0]
        parent = data['matedEntities'][1]['matedOccurrence'][0]

        if data['name'][0:3] == 'dof':
            parts = data['name'].split('_')
            del parts[0]
            data['inverted'] = False
            if parts[-1] == 'inv' or parts[-1] == 'inverted':
                data['inverted'] = True
                del parts[-1]
            name = '_'.join(parts)
            if name == '':
                print(Fore.RED +
                      'ERROR: a DOF dones\'t have any name ("'+data['name']+'" should be "dof_...")' + Style.RESET_ALL)
                exit()

            if data['mateType'] == 'REVOLUTE' or data['mateType'] == 'CYLINDRICAL':
                jointType = 'revolute'
                limits = getLimits(jointType, data['name'])
            elif data['mateType'] == 'SLIDER':
                jointType = 'prismatic'
                limits = getLimits(jointType, data['name'])
            elif data['mateType'] == 'FASTENED':
                jointType = 'floating'
                limits = None
            else:
                print(Fore.RED + 'ERROR: "' + name +
                      '" is declared as a DOF but the mate type is '+data['mateType']+'')
                print(
                    '       Only REVOLUTE, CYLINDRICAL, SLIDER and FASTENED are supported' + Style.RESET_ALL)
                exit(1)

            # We compute the axis in the world frame
            matedEntity = data['matedEntities'][0]
            matedTransform = getOccurrence(
                matedEntity['matedOccurrence'])['transform']

            # jointToPart is the (rotation only) matrix from joint to the part
            # it is attached to
            jointToPart = np.eye(4)
            jointToPart[:3, :3] = np.stack((
                np.array(matedEntity['matedCS']['xAxis']),
                np.array(matedEntity['matedCS']['yAxis']),
                np.array(matedEntity['matedCS']['zAxis'])
            ))

            if data['inverted']:
                if limits is not None:
                    limits = (-limits[1], -limits[0])

                # Flipping the joint around X axis
                flip = np.array([[1, 0, 0, 0],
                                [0, -1, 0, 0],
                                [0, 0, -1, 0],
                                [0, 0,  0,  1]])
                jointToPart = flip.dot(jointToPart)

            zAxis = np.array([0, 0, 1])

            origin = matedEntity['matedCS']['origin']
            translation = np.matrix(np.identity(4))
            translation[0, 3] += origin[0]
            translation[1, 3] += origin[1]
            translation[2, 3] += origin[2]
            worldAxisFrame = matedTransform * translation

            # Resulting frame of axis, always revolving around z
            worldAxisFrame = worldAxisFrame.dot(jointBase)

            limitsStr = ''
            if limits is not None:
                limitsStr = '[' + str(round(limits[0], 3)) + \
                    ': ' + str(round(limits[1], 3)) + ']'
            print(Fore.GREEN + '+ Found DOF: '+name + ' ' + Style.DIM +
                  '('+jointType+')'+limitsStr + Style.RESET_ALL)

            if child in relations:
                print(Fore.RED)
                print('Error, the relation '+name +
                      ' is connected a child that is already connected')
                print('Be sure you ordered properly your relations, see:')
                print(
                    'https://onshape-to-robot.readthedocs.io/en/latest/design.html#specifying-degrees-of-freedom')
                print(Style.RESET_ALL)
                exit()

            relations[child] = {
                'parent': parent,
                'worldAxisFrame': worldAxisFrame,
                'zAxis': zAxis,
                'name': name,
                'type': jointType,
                'limits': limits
            }

            assignParts(child, child)
            assignParts(parent, parent)
            if child not in frames:
                frames[child] = []
            if parent not in frames:
                frames[parent] = []

print(Fore.GREEN + Style.BRIGHT + '* Found total ' +
      str(len(relations))+' DOFs' + Style.RESET_ALL)

# If we have no DOF
if len(relations) == 0:
    trunk = root['instances'][0]['id']
    assignParts(trunk, trunk)

def connectParts(child, parent):
    assignParts(child, parent)

# Spreading parts assignations, this parts mainly does two things:
# 1. Finds the parts of the top level assembly that are not directly in a sub assembly and try to assign them
#    to an existing link that was identified before
# 2. Among those parts, finds the ones that are frames (connected with a frame_* connector)
changed = True
while changed:
    changed = False
    for feature in features:
        if feature['featureType'] != 'mate' or feature['suppressed']:
            continue

        data = feature['featureData']

        if len(data['matedEntities']) != 2 \
                or len(data['matedEntities'][0]['matedOccurrence']) == 0 \
                or len(data['matedEntities'][1]['matedOccurrence']) == 0:
            continue

        occurrenceA = data['matedEntities'][0]['matedOccurrence'][0]
        occurrenceB = data['matedEntities'][1]['matedOccurrence'][0]

        if (occurrenceA not in assignations) != (occurrenceB not in assignations):
            if data['name'][0:5] == 'frame':
                name = '_'.join(data['name'].split('_')[1:])
                if occurrenceA in assignations:
                    frames[occurrenceA].append(
                        [name, data['matedEntities'][1]['matedOccurrence']])
                    assignParts(occurrenceB, {True: assignations[occurrenceA], False: 'frame'}[
                                config['drawFrames']])
                    changed = True
                else:
                    frames[occurrenceB].append(
                        [name, data['matedEntities'][0]['matedOccurrence']])
                    assignParts(occurrenceA, {True: assignations[occurrenceB], False: 'frame'}[
                                config['drawFrames']])
                    changed = True
            else:
                if occurrenceA in assignations:
                    connectParts(occurrenceB, assignations[occurrenceA])
                    changed = True
                else:
                    connectParts(occurrenceA, assignations[occurrenceB])
                    changed = True

# Building and checking robot tree, here we:
# 1. Search for robot trunk (which will be the top-level link)
# 2. Scan for orphaned parts (if you add something floating with no mate to anything)
#    that are then assigned to trunk by default
# 3. Collect all the pieces of the robot tree
print("\n" + Style.BRIGHT + '* Building robot tree' + Style.RESET_ALL)

for childId in relations:
    entry = relations[childId]
    if entry['parent'] not in relations:
        trunk = entry['parent']
        break
trunkOccurrence = getOccurrence([trunk])
print(Style.BRIGHT + '* Trunk is ' +
      trunkOccurrence['instance']['name'] + Style.RESET_ALL)

for occurrence in occurrences.values():
    if occurrence['assignation'] is None:
        print(Fore.YELLOW + 'WARNING: part (' +
              occurrence['instance']['name']+') has no assignation, connecting it with trunk' + Style.RESET_ALL)
        child = occurrence['path'][0]
        connectParts(child, trunk)

def collect(id):
    part = {}
    part['id'] = id
    part['children'] = []
    for childId in relations:
        entry = relations[childId]
        if entry['parent'] == id:
            child = collect(childId)
            child['axis_frame'] = entry['worldAxisFrame']
            child['z_axis'] = entry['zAxis']
            child['dof_name'] = entry['name']
            child['jointType'] = entry['type']
            child['jointLimits'] = entry['limits']
            part['children'].append(child)
    return part

tree = collect(trunk)
<commit_msg>Fixing typo<commit_after>import math
from sys import exit
import numpy as np
import uuid
from .onshape_api.client import Client
from .config import config, configFile
from colorama import Fore, Back, Style

# OnShape API client
client = Client(logging=False, creds=configFile)
client.useCollisionsConfigurations = config['useCollisionsConfigurations']

# If a versionId is provided, it will be used, else the main workspace is retrieved
if config['versionId'] == '':
    print("\n" + Style.BRIGHT + '* Retrieving workspace ID ...' + Style.RESET_ALL)
    document = client.get_document(config['documentId']).json()
    workspaceId = document['defaultWorkspace']['id']
    print(Fore.GREEN + "+ Using workspace id: " + workspaceId + Style.RESET_ALL)
else:
    print("\n" + Style.BRIGHT + '* Using configuration version ID ' +
          config['versionId']+' ...' + Style.RESET_ALL)

# Now, finding the assembly, according to given name in configuration, or else the first possible one
print("\n" + Style.BRIGHT +
      '* Retrieving elements in the document, searching for the assembly...' + Style.RESET_ALL)
if config['versionId'] == '':
    elements = client.list_elements(config['documentId'], workspaceId).json()
else:
    elements = client.list_elements(
        config['documentId'], config['versionId'], 'v').json()
assemblyId = None
assemblyName = ''
for element in elements:
    if element['type'] == 'Assembly' and \
            (config['assemblyName'] is False or element['name'] == config['assemblyName']):
        print(Fore.GREEN + "+ Found assembly, id: " +
              element['id']+', name: "'+element['name']+'"' + Style.RESET_ALL)
        assemblyName = element['name']
        assemblyId = element['id']

if assemblyId == None:
    print(Fore.RED + "ERROR: Unable to find assembly in this document" + Style.RESET_ALL)
    exit(1)

# Retrieving the assembly
print("\n" + Style.BRIGHT + '* Retrieving assembly "' +
      assemblyName+'" with id '+assemblyId + Style.RESET_ALL)
if config['versionId'] == '':
    assembly = client.get_assembly(
        config['documentId'], workspaceId, assemblyId)
else:
    assembly = client.get_assembly(
        config['documentId'], config['versionId'], assemblyId, 'v')

root = assembly['rootAssembly']

# Finds a (leaf) instance given the full path, typically A B C where A and B would be subassemblies and C
# the final part

def findInstance(path, instances=None):
    global assembly

    if instances is None:
        instances = assembly['rootAssembly']['instances']

    for instance in instances:
        if instance['id'] == path[0]:
            if len(path) == 1:
                # If the length of remaining path is 1, the part is in the current assembly/subassembly
                return instance
            else:
                # Else, we need to find the matching sub assembly to find the proper part (recursively)
                d = instance['documentId']
                m = instance['documentMicroversion']
                e = instance['elementId']
                for asm in assembly['subAssemblies']:
                    if asm['documentId'] == d and asm['documentMicroversion'] == m and asm['elementId'] == e:
                        return findInstance(path[1:], asm['instances'])

    print(Fore.RED + 'Could not find instance for ' + str(path) + Style.RESET_ALL)

# Collecting occurrences, the path is the assembly / sub assembly chain
occurrences = {}
for occurrence in root['occurrences']:
    occurrence['assignation'] = None
    occurrence['instance'] = findInstance(occurrence['path'])
    occurrence['transform'] = np.matrix(
        np.reshape(occurrence['transform'], (4, 4)))
    occurrence['linkName'] = None
    occurrences[tuple(occurrence['path'])] = occurrence

# Gets an occurrence given its full path

def getOccurrence(path):
    return occurrences[tuple(path)]

# Assignations are pieces that will be in the same link. Note that this is only for top-level
# item of the path (all sub assemblies and parts in assemblies are naturally in the same link as
# the parent), but other parts that can be connected with mates in top assemblies are then assigned to
# the link
assignations = {}

# Frames (mated with frame_ name) will be special links in the output file allowing to track some specific
# manually identified frames
frames = {}

def assignParts(root, parent):
    assignations[root] = parent
    for occurrence in occurrences.values():
        if occurrence['path'][0] == root:
            occurrence['assignation'] = parent

from .features import init as features_init, getLimits
features_init(client, config, root, workspaceId, assemblyId)

# First, features are scanned to find the DOFs. Links that they connects are then tagged
print("\n" + Style.BRIGHT +
      '* Getting assembly features, scanning for DOFs...' + Style.RESET_ALL)
trunk = None
relations = {}
features = root['features']
for feature in features:
    if feature['featureType'] == 'mateConnector':
        name = feature['featureData']['name']
        if name[0:5] == 'link_':
            name = name[5:]
            occurrences[(feature['featureData']['occurrence'][0],)
                        ]['linkName'] = name
    else:
        if feature['suppressed']:
            continue

        data = feature['featureData']

        if len(data['matedEntities']) != 2 or \
                len(data['matedEntities'][0]['matedOccurrence']) == 0 \
                or len(data['matedEntities'][1]['matedOccurrence']) == 0:
            continue

        child = data['matedEntities'][0]['matedOccurrence'][0]
        parent = data['matedEntities'][1]['matedOccurrence'][0]

        if data['name'][0:3] == 'dof':
            parts = data['name'].split('_')
            del parts[0]
            data['inverted'] = False
            if parts[-1] == 'inv' or parts[-1] == 'inverted':
                data['inverted'] = True
                del parts[-1]
            name = '_'.join(parts)
            if name == '':
                print(Fore.RED +
                      'ERROR: a DOF dones\'t have any name ("'+data['name']+'" should be "dof_...")' + Style.RESET_ALL)
                exit()

            if data['mateType'] == 'REVOLUTE' or data['mateType'] == 'CYLINDRICAL':
                jointType = 'revolute'
                limits = getLimits(jointType, data['name'])
            elif data['mateType'] == 'SLIDER':
                jointType = 'prismatic'
                limits = getLimits(jointType, data['name'])
            elif data['mateType'] == 'FASTENED':
                jointType = 'floating'
                limits = None
            else:
                print(Fore.RED + 'ERROR: "' + name +
                      '" is declared as a DOF but the mate type is '+data['mateType']+'')
                print(
                    '       Only REVOLUTE, CYLINDRICAL, SLIDER and FASTENED are supported' + Style.RESET_ALL)
                exit(1)

            # We compute the axis in the world frame
            matedEntity = data['matedEntities'][0]
            matedTransform = getOccurrence(
                matedEntity['matedOccurrence'])['transform']

            # jointToPart is the (rotation only) matrix from joint to the part
            # it is attached to
            jointToPart = np.eye(4)
            jointToPart[:3, :3] = np.stack((
                np.array(matedEntity['matedCS']['xAxis']),
                np.array(matedEntity['matedCS']['yAxis']),
                np.array(matedEntity['matedCS']['zAxis'])
            ))

            if data['inverted']:
                if limits is not None:
                    limits = (-limits[1], -limits[0])

                # Flipping the joint around X axis
                flip = np.array([[1, 0, 0, 0],
                                [0, -1, 0, 0],
                                [0, 0, -1, 0],
                                [0, 0,  0,  1]])
                jointToPart = flip.dot(jointToPart)

            zAxis = np.array([0, 0, 1])

            origin = matedEntity['matedCS']['origin']
            translation = np.matrix(np.identity(4))
            translation[0, 3] += origin[0]
            translation[1, 3] += origin[1]
            translation[2, 3] += origin[2]
            worldAxisFrame = matedTransform * translation

            # Resulting frame of axis, always revolving around z
            worldAxisFrame = worldAxisFrame.dot(jointToPart)

            limitsStr = ''
            if limits is not None:
                limitsStr = '[' + str(round(limits[0], 3)) + \
                    ': ' + str(round(limits[1], 3)) + ']'
            print(Fore.GREEN + '+ Found DOF: '+name + ' ' + Style.DIM +
                  '('+jointType+')'+limitsStr + Style.RESET_ALL)

            if child in relations:
                print(Fore.RED)
                print('Error, the relation '+name +
                      ' is connected a child that is already connected')
                print('Be sure you ordered properly your relations, see:')
                print(
                    'https://onshape-to-robot.readthedocs.io/en/latest/design.html#specifying-degrees-of-freedom')
                print(Style.RESET_ALL)
                exit()

            relations[child] = {
                'parent': parent,
                'worldAxisFrame': worldAxisFrame,
                'zAxis': zAxis,
                'name': name,
                'type': jointType,
                'limits': limits
            }

            assignParts(child, child)
            assignParts(parent, parent)
            if child not in frames:
                frames[child] = []
            if parent not in frames:
                frames[parent] = []

print(Fore.GREEN + Style.BRIGHT + '* Found total ' +
      str(len(relations))+' DOFs' + Style.RESET_ALL)

# If we have no DOF
if len(relations) == 0:
    trunk = root['instances'][0]['id']
    assignParts(trunk, trunk)

def connectParts(child, parent):
    assignParts(child, parent)

# Spreading parts assignations, this parts mainly does two things:
# 1. Finds the parts of the top level assembly that are not directly in a sub assembly and try to assign them
#    to an existing link that was identified before
# 2. Among those parts, finds the ones that are frames (connected with a frame_* connector)
changed = True
while changed:
    changed = False
    for feature in features:
        if feature['featureType'] != 'mate' or feature['suppressed']:
            continue

        data = feature['featureData']

        if len(data['matedEntities']) != 2 \
                or len(data['matedEntities'][0]['matedOccurrence']) == 0 \
                or len(data['matedEntities'][1]['matedOccurrence']) == 0:
            continue

        occurrenceA = data['matedEntities'][0]['matedOccurrence'][0]
        occurrenceB = data['matedEntities'][1]['matedOccurrence'][0]

        if (occurrenceA not in assignations) != (occurrenceB not in assignations):
            if data['name'][0:5] == 'frame':
                name = '_'.join(data['name'].split('_')[1:])
                if occurrenceA in assignations:
                    frames[occurrenceA].append(
                        [name, data['matedEntities'][1]['matedOccurrence']])
                    assignParts(occurrenceB, {True: assignations[occurrenceA], False: 'frame'}[
                                config['drawFrames']])
                    changed = True
                else:
                    frames[occurrenceB].append(
                        [name, data['matedEntities'][0]['matedOccurrence']])
                    assignParts(occurrenceA, {True: assignations[occurrenceB], False: 'frame'}[
                                config['drawFrames']])
                    changed = True
            else:
                if occurrenceA in assignations:
                    connectParts(occurrenceB, assignations[occurrenceA])
                    changed = True
                else:
                    connectParts(occurrenceA, assignations[occurrenceB])
                    changed = True

# Building and checking robot tree, here we:
# 1. Search for robot trunk (which will be the top-level link)
# 2. Scan for orphaned parts (if you add something floating with no mate to anything)
#    that are then assigned to trunk by default
# 3. Collect all the pieces of the robot tree
print("\n" + Style.BRIGHT + '* Building robot tree' + Style.RESET_ALL)

for childId in relations:
    entry = relations[childId]
    if entry['parent'] not in relations:
        trunk = entry['parent']
        break
trunkOccurrence = getOccurrence([trunk])
print(Style.BRIGHT + '* Trunk is ' +
      trunkOccurrence['instance']['name'] + Style.RESET_ALL)

for occurrence in occurrences.values():
    if occurrence['assignation'] is None:
        print(Fore.YELLOW + 'WARNING: part (' +
              occurrence['instance']['name']+') has no assignation, connecting it with trunk' + Style.RESET_ALL)
        child = occurrence['path'][0]
        connectParts(child, trunk)

def collect(id):
    part = {}
    part['id'] = id
    part['children'] = []
    for childId in relations:
        entry = relations[childId]
        if entry['parent'] == id:
            child = collect(childId)
            child['axis_frame'] = entry['worldAxisFrame']
            child['z_axis'] = entry['zAxis']
            child['dof_name'] = entry['name']
            child['jointType'] = entry['type']
            child['jointLimits'] = entry['limits']
            part['children'].append(child)
    return part

tree = collect(trunk)
<|endoftext|>