<commit_before># -*- coding: utf-8 -*-

from __future__ import absolute_import

# Import Python libs
import contextlib
import errno
import logging
import os
import re
import shutil
import subprocess
import time
import urllib

# Import salt libs
import salt.utils
import salt.modules.selinux
from salt.exceptions import CommandExecutionError, FileLockError, MinionError

# Import 3rd-party libs
from salt.ext import six

log = logging.getLogger(__name__)

LOCAL_PROTOS = ('', 'file')
REMOTE_PROTOS = ('http', 'https', 'ftp', 'swift', 's3')
VALID_PROTOS = ('salt', 'file') + REMOTE_PROTOS
TEMPFILE_PREFIX = '__salt.tmp.'

HASHES = {
    'sha512': 128,
    'sha384': 96,
    'sha256': 64,
    'sha224': 56,
    'sha1': 40,
    'md5': 32,
}
HASHES_REVMAP = dict([(y, x) for x, y in six.iteritems(HASHES)])

def guess_archive_type(name):
    '''
    Guess an archive type (tar, zip, or rar) by its file extension
    '''
    name = name.lower()
    for ending in ('tar', 'tar.gz', 'tar.bz2', 'tar.xz', 'tgz', 'tbz2', 'txz',
                   'tar.lzma', 'tlz'):
        if name.endswith('.' + ending):
            return 'tar'
    for ending in ('zip', 'rar'):
        if name.endswith('.' + ending):
            return ending
    return None

def mkstemp(*args, **kwargs):
    '''
    Should eventually reside here, but for now point back at old location in
    salt.utils
    '''
    return salt.utils.mkstemp(*args, **kwargs)

def recursive_copy(source, dest):
    '''
    Recursively copy the source directory to the destination,
    leaving files with the source does not explicitly overwrite.

    (identical to cp -r on a unix machine)
    '''
    for root, _, files in os.walk(source):
        path_from_source = root.replace(source, '').lstrip(os.sep)
        target_directory = os.path.join(dest, path_from_source)
        if not os.path.exists(target_directory):
            os.makedirs(target_directory)
        for name in files:
            file_path_from_source = os.path.join(source, path_from_source, name)
            target_path = os.path.join(target_directory, name)
            shutil.copyfile(file_path_from_source, target_path)

def copyfile(source, dest, backup_mode='', cachedir=''):
    '''
    Copy files from a source to a destination in an atomic way, and if
    specified cache the file.
    '''
    if not os.path.isfile(source):
        raise IOError(
            '[Errno 2] No such file or directory: {0}'.format(source)
        )
    if not os.path.isdir(os.path.dirname(dest)):
        raise IOError(
            '[Errno 2] No such file or directory: {0}'.format(dest)
        )
    bname = os.path.basename(dest)
    dname = os.path.dirname(os.path.abspath(dest))
    tgt = mkstemp(prefix=bname, dir=dname)
    shutil.copyfile(source, tgt)
    bkroot = ''
    if cachedir:
        bkroot = os.path.join(cachedir, 'file_backup')
    if backup_mode == 'minion' or backup_mode == 'both' and bkroot:
        if os.path.exists(dest):
            salt.utils.backup_minion(dest, bkroot)
    if backup_mode == 'master' or backup_mode == 'both' and bkroot:
        # TODO, backup to master
        pass
    # Get current file stats to they can be replicated after the new file is
    # moved to the destination path.
    fstat = None
    if not salt.utils.is_windows():
        try:
            fstat = os.stat(dest)
        except OSError:
            pass
    shutil.move(tgt, dest)
    if fstat is not None:
        os.chown(dest, fstat.st_uid, fstat.st_gid)
        os.chmod(dest, fstat.st_mode)
    # If SELINUX is available run a restorecon on the file
    rcon = salt.utils.which('restorecon')
    if rcon:
        policy = False
        try:
            policy = salt.modules.selinux.getenforce()
        except (ImportError, CommandExecutionError):
            pass
        if policy == 'Enforcing':
            with salt.utils.fopen(os.devnull, 'w') as dev_null:
                cmd = [rcon, dest]
                subprocess.call(cmd, stdout=dev_null, stderr=dev_null)
    if os.path.isfile(tgt):
        # The temp file failed to move
        try:
            os.remove(tgt)
        except Exception:
            pass

def rename(src, dst):
    '''
    On Windows, os.rename() will fail with a WindowsError exception if a file
    exists at the destination path. This function checks for this error and if
    found, it deletes the destination path first.
    '''
    try:
        os.rename(src, dst)
    except OSError as exc:
        if exc.errno != errno.EEXIST:
            raise
        try:
            os.remove(dst)
        except OSError as exc:
            if exc.errno != errno.ENOENT:
                raise MinionError(
                    'Error: Unable to remove {0}: {1}'.format(
                        dst,
                        exc.strerror
                    )
                )
        os.rename(src, dst)

def process_read_exception(exc, path):
    '''
    Common code for raising exceptions when reading a file fails
    '''
    if exc.errno == errno.ENOENT:
        raise CommandExecutionError('{0} does not exist'.format(path))
    elif exc.errno == errno.EACCES:
        raise CommandExecutionError(
            'Permission denied reading from {0}'.format(path)
        )
    else:
        raise CommandExecutionError(
            'Error {0} encountered reading from {1}: {2}'.format(
                exc.errno, path, exc.strerror
            )
        )

@contextlib.contextmanager
def wait_lock(path, lock_fn=None, timeout=5, sleep=0.1, time_start=None):
    '''
    Obtain a write lock. If one exists, wait for it to release first
    '''
    if not isinstance(path, six.string_types):
        raise FileLockError('path must be a string')
    if lock_fn is None:
        lock_fn = path + '.w'
    if time_start is None:
        time_start = time.time()
    obtained_lock = False

    def _raise_error(msg, race=False):
        '''
        Raise a FileLockError
        '''
        raise FileLockError(msg, time_start=time_start)

    try:
        if os.path.exists(lock_fn) and not os.path.isfile(lock_fn):
            _raise_error(
                'lock_fn {0} exists and is not a file'.format(lock_fn)
            )

        open_flags = os.O_CREAT | os.O_EXCL | os.O_WRONLY
        while time.time() - time_start < timeout:
            try:
                # Use os.open() to obtain filehandle so that we can force an
                # exception if the file already exists. Concept found here:
                # http://stackoverflow.com/a/10979569
                fh_ = os.open(lock_fn, open_flags)
            except (IOError, OSError) as exc:
                if exc.errno != errno.EEXIST:
                    _raise_error(
                        'Error {0} encountered obtaining file lock {1}: {2}'
                        .format(exc.errno, lock_fn, exc.strerror)
                    )
                log.trace(
                    'Lock file %s exists, sleeping %f seconds', lock_fn, sleep
                )
                time.sleep(sleep)
            else:
                # Write the lock file
                with os.fdopen(fh_, 'w'):
                    pass
                # Lock successfully acquired
                log.trace('Write lock %s obtained', lock_fn)
                obtained_lock = True
                # Transfer control back to the code inside the with block
                yield
                # Exit the loop
                break

        else:
            _raise_error(
                'Timeout of {0} seconds exceeded waiting for lock_fn {1} '
                'to be released'.format(timeout, lock_fn)
            )

    except FileLockError:
        raise

    except Exception as exc:
        _raise_error(
            'Error encountered obtaining file lock {0}: {1}'.format(
                lock_fn,
                exc
            )
        )

    finally:
        if obtained_lock:
            os.remove(lock_fn)
            log.trace('Write lock for %s (%s) released', path, lock_fn)

@contextlib.contextmanager
def set_umask(mask):
    '''
    Temporarily set the umask and restore once the contextmanager exits
    '''
    if salt.utils.is_windows():
        # Don't attempt on Windows
        yield
    else:
        try:
            orig_mask = os.umask(mask)
            yield
        finally:
            os.umask(orig_mask)

def safe_filename_leaf(file_basename):
    '''
    Input the basename of a file, without the directory tree, and returns a safe name to use
    i.e. only the required characters are converted by urllib.quote
    If the input is a PY2 String, output a PY2 String. If input is Unicode output Unicode.
    For consistency all platforms are treated the same. Hard coded to utf8 as its ascii compatible
    windows is \\ / : * ? " < > | posix is /

    .. versionadded:: 2017.7.2

    :codeauthor: Damon Atkins <https://github.com/damon-atkins>
    '''
    def _replace(re_obj):
        return urllib.quote(re_obj.group(0), safe=u'')
    if not isinstance(file_basename, six.text_type):
        # the following string is not prefixed with u
        return re.sub('[\\\\:/*?"<>|]',
                      _replace,
                      six.text_type(file_basename, 'utf8').encode('ascii', 'backslashreplace'))
    # the following string is prefixed with u
    return re.sub(u'[\\\\:/*?"<>|]', _replace, file_basename, flags=re.UNICODE)

def safe_filepath(file_path_name, dir_sep=None):
    '''
    Input the full path and filename, splits on directory separator and calls safe_filename_leaf for
    each part of the path. dir_sep allows coder to force a directory separate to a particular character

    .. versionadded:: 2017.7.2

    :codeauthor: Damon Atkins <https://github.com/damon-atkins>
    '''
    if not dir_sep:
        dir_sep = os.sep
    # Normally if file_path_name or dir_sep is Unicode then the output will be Unicode
    # This code ensure the output type is the same as file_path_name
    if not isinstance(file_path_name, six.text_type) and isinstance(dir_sep, six.text_type):
        dir_sep = dir_sep.encode('ascii')  # This should not be executed under PY3
    # splitdrive only set drive on windows platform
    (drive, path) = os.path.splitdrive(file_path_name)
    path = dir_sep.join([safe_filename_leaf(file_section) for file_section in path.rsplit(dir_sep)])
    if drive:
        path = dir_sep.join([drive, path])
    return path

def remove(path):
    '''
    Runs os.remove(path) and suppresses the OSError if the file doesn't exist
    '''
    try:
        os.remove(path)
    except OSError as exc:
        if exc.errno != errno.ENOENT:
            raise

<commit_msg>remove blank line at end of file<commit_after># -*- coding: utf-8 -*-

from __future__ import absolute_import

# Import Python libs
import contextlib
import errno
import logging
import os
import re
import shutil
import subprocess
import time
import urllib

# Import salt libs
import salt.utils
import salt.modules.selinux
from salt.exceptions import CommandExecutionError, FileLockError, MinionError

# Import 3rd-party libs
from salt.ext import six

log = logging.getLogger(__name__)

LOCAL_PROTOS = ('', 'file')
REMOTE_PROTOS = ('http', 'https', 'ftp', 'swift', 's3')
VALID_PROTOS = ('salt', 'file') + REMOTE_PROTOS
TEMPFILE_PREFIX = '__salt.tmp.'

HASHES = {
    'sha512': 128,
    'sha384': 96,
    'sha256': 64,
    'sha224': 56,
    'sha1': 40,
    'md5': 32,
}
HASHES_REVMAP = dict([(y, x) for x, y in six.iteritems(HASHES)])

def guess_archive_type(name):
    '''
    Guess an archive type (tar, zip, or rar) by its file extension
    '''
    name = name.lower()
    for ending in ('tar', 'tar.gz', 'tar.bz2', 'tar.xz', 'tgz', 'tbz2', 'txz',
                   'tar.lzma', 'tlz'):
        if name.endswith('.' + ending):
            return 'tar'
    for ending in ('zip', 'rar'):
        if name.endswith('.' + ending):
            return ending
    return None

def mkstemp(*args, **kwargs):
    '''
    Should eventually reside here, but for now point back at old location in
    salt.utils
    '''
    return salt.utils.mkstemp(*args, **kwargs)

def recursive_copy(source, dest):
    '''
    Recursively copy the source directory to the destination,
    leaving files with the source does not explicitly overwrite.

    (identical to cp -r on a unix machine)
    '''
    for root, _, files in os.walk(source):
        path_from_source = root.replace(source, '').lstrip(os.sep)
        target_directory = os.path.join(dest, path_from_source)
        if not os.path.exists(target_directory):
            os.makedirs(target_directory)
        for name in files:
            file_path_from_source = os.path.join(source, path_from_source, name)
            target_path = os.path.join(target_directory, name)
            shutil.copyfile(file_path_from_source, target_path)

def copyfile(source, dest, backup_mode='', cachedir=''):
    '''
    Copy files from a source to a destination in an atomic way, and if
    specified cache the file.
    '''
    if not os.path.isfile(source):
        raise IOError(
            '[Errno 2] No such file or directory: {0}'.format(source)
        )
    if not os.path.isdir(os.path.dirname(dest)):
        raise IOError(
            '[Errno 2] No such file or directory: {0}'.format(dest)
        )
    bname = os.path.basename(dest)
    dname = os.path.dirname(os.path.abspath(dest))
    tgt = mkstemp(prefix=bname, dir=dname)
    shutil.copyfile(source, tgt)
    bkroot = ''
    if cachedir:
        bkroot = os.path.join(cachedir, 'file_backup')
    if backup_mode == 'minion' or backup_mode == 'both' and bkroot:
        if os.path.exists(dest):
            salt.utils.backup_minion(dest, bkroot)
    if backup_mode == 'master' or backup_mode == 'both' and bkroot:
        # TODO, backup to master
        pass
    # Get current file stats to they can be replicated after the new file is
    # moved to the destination path.
    fstat = None
    if not salt.utils.is_windows():
        try:
            fstat = os.stat(dest)
        except OSError:
            pass
    shutil.move(tgt, dest)
    if fstat is not None:
        os.chown(dest, fstat.st_uid, fstat.st_gid)
        os.chmod(dest, fstat.st_mode)
    # If SELINUX is available run a restorecon on the file
    rcon = salt.utils.which('restorecon')
    if rcon:
        policy = False
        try:
            policy = salt.modules.selinux.getenforce()
        except (ImportError, CommandExecutionError):
            pass
        if policy == 'Enforcing':
            with salt.utils.fopen(os.devnull, 'w') as dev_null:
                cmd = [rcon, dest]
                subprocess.call(cmd, stdout=dev_null, stderr=dev_null)
    if os.path.isfile(tgt):
        # The temp file failed to move
        try:
            os.remove(tgt)
        except Exception:
            pass

def rename(src, dst):
    '''
    On Windows, os.rename() will fail with a WindowsError exception if a file
    exists at the destination path. This function checks for this error and if
    found, it deletes the destination path first.
    '''
    try:
        os.rename(src, dst)
    except OSError as exc:
        if exc.errno != errno.EEXIST:
            raise
        try:
            os.remove(dst)
        except OSError as exc:
            if exc.errno != errno.ENOENT:
                raise MinionError(
                    'Error: Unable to remove {0}: {1}'.format(
                        dst,
                        exc.strerror
                    )
                )
        os.rename(src, dst)

def process_read_exception(exc, path):
    '''
    Common code for raising exceptions when reading a file fails
    '''
    if exc.errno == errno.ENOENT:
        raise CommandExecutionError('{0} does not exist'.format(path))
    elif exc.errno == errno.EACCES:
        raise CommandExecutionError(
            'Permission denied reading from {0}'.format(path)
        )
    else:
        raise CommandExecutionError(
            'Error {0} encountered reading from {1}: {2}'.format(
                exc.errno, path, exc.strerror
            )
        )

@contextlib.contextmanager
def wait_lock(path, lock_fn=None, timeout=5, sleep=0.1, time_start=None):
    '''
    Obtain a write lock. If one exists, wait for it to release first
    '''
    if not isinstance(path, six.string_types):
        raise FileLockError('path must be a string')
    if lock_fn is None:
        lock_fn = path + '.w'
    if time_start is None:
        time_start = time.time()
    obtained_lock = False

    def _raise_error(msg, race=False):
        '''
        Raise a FileLockError
        '''
        raise FileLockError(msg, time_start=time_start)

    try:
        if os.path.exists(lock_fn) and not os.path.isfile(lock_fn):
            _raise_error(
                'lock_fn {0} exists and is not a file'.format(lock_fn)
            )

        open_flags = os.O_CREAT | os.O_EXCL | os.O_WRONLY
        while time.time() - time_start < timeout:
            try:
                # Use os.open() to obtain filehandle so that we can force an
                # exception if the file already exists. Concept found here:
                # http://stackoverflow.com/a/10979569
                fh_ = os.open(lock_fn, open_flags)
            except (IOError, OSError) as exc:
                if exc.errno != errno.EEXIST:
                    _raise_error(
                        'Error {0} encountered obtaining file lock {1}: {2}'
                        .format(exc.errno, lock_fn, exc.strerror)
                    )
                log.trace(
                    'Lock file %s exists, sleeping %f seconds', lock_fn, sleep
                )
                time.sleep(sleep)
            else:
                # Write the lock file
                with os.fdopen(fh_, 'w'):
                    pass
                # Lock successfully acquired
                log.trace('Write lock %s obtained', lock_fn)
                obtained_lock = True
                # Transfer control back to the code inside the with block
                yield
                # Exit the loop
                break

        else:
            _raise_error(
                'Timeout of {0} seconds exceeded waiting for lock_fn {1} '
                'to be released'.format(timeout, lock_fn)
            )

    except FileLockError:
        raise

    except Exception as exc:
        _raise_error(
            'Error encountered obtaining file lock {0}: {1}'.format(
                lock_fn,
                exc
            )
        )

    finally:
        if obtained_lock:
            os.remove(lock_fn)
            log.trace('Write lock for %s (%s) released', path, lock_fn)

@contextlib.contextmanager
def set_umask(mask):
    '''
    Temporarily set the umask and restore once the contextmanager exits
    '''
    if salt.utils.is_windows():
        # Don't attempt on Windows
        yield
    else:
        try:
            orig_mask = os.umask(mask)
            yield
        finally:
            os.umask(orig_mask)

def safe_filename_leaf(file_basename):
    '''
    Input the basename of a file, without the directory tree, and returns a safe name to use
    i.e. only the required characters are converted by urllib.quote
    If the input is a PY2 String, output a PY2 String. If input is Unicode output Unicode.
    For consistency all platforms are treated the same. Hard coded to utf8 as its ascii compatible
    windows is \\ / : * ? " < > | posix is /

    .. versionadded:: 2017.7.2

    :codeauthor: Damon Atkins <https://github.com/damon-atkins>
    '''
    def _replace(re_obj):
        return urllib.quote(re_obj.group(0), safe=u'')
    if not isinstance(file_basename, six.text_type):
        # the following string is not prefixed with u
        return re.sub('[\\\\:/*?"<>|]',
                      _replace,
                      six.text_type(file_basename, 'utf8').encode('ascii', 'backslashreplace'))
    # the following string is prefixed with u
    return re.sub(u'[\\\\:/*?"<>|]', _replace, file_basename, flags=re.UNICODE)

def safe_filepath(file_path_name, dir_sep=None):
    '''
    Input the full path and filename, splits on directory separator and calls safe_filename_leaf for
    each part of the path. dir_sep allows coder to force a directory separate to a particular character

    .. versionadded:: 2017.7.2

    :codeauthor: Damon Atkins <https://github.com/damon-atkins>
    '''
    if not dir_sep:
        dir_sep = os.sep
    # Normally if file_path_name or dir_sep is Unicode then the output will be Unicode
    # This code ensure the output type is the same as file_path_name
    if not isinstance(file_path_name, six.text_type) and isinstance(dir_sep, six.text_type):
        dir_sep = dir_sep.encode('ascii')  # This should not be executed under PY3
    # splitdrive only set drive on windows platform
    (drive, path) = os.path.splitdrive(file_path_name)
    path = dir_sep.join([safe_filename_leaf(file_section) for file_section in path.rsplit(dir_sep)])
    if drive:
        path = dir_sep.join([drive, path])
    return path

def remove(path):
    '''
    Runs os.remove(path) and suppresses the OSError if the file doesn't exist
    '''
    try:
        os.remove(path)
    except OSError as exc:
        if exc.errno != errno.ENOENT:
            raise
<|endoftext|>


<commit_before>'''
Defines a ServerGroup class for running and managing servers, as well as a set of
base server types.

The following information is needed to start a server:
1- handler / wsgi app
2- address
3- ssl preference

'''
import collections
import socket
import os
import threading

if hasattr(os, "fork"):
    import ufork
else:
    ufork = None
from gevent.pywsgi import WSGIServer
import gevent.socket

import async
from protected import Protected
import context

class ServerGroup(object):
    def __init__(self, wsgi_apps=(), stream_handlers=(), prefork=False, daemonize=False, dev=False, **kw):
        self.prefork = prefork
        self.daemonize = daemonize
        self.post_fork = kw.get('post_fork')  # callback to be executed post fork
        self.server_log = kw.get('gevent_log')
        self.wsgi_apps = wsgi_apps
        self.stream_handlers = stream_handlers
        ctx = context.get_context()
        self.num_workers = ctx.num_workers
        self.servers = []
        self.socks = {}
        for app in wsgi_apps:
            app, address, ssl = app
            sock = _make_server_sock(address)
            if isinstance(ssl, Protected):
                protected = ssl
            elif ssl:
                protected = ctx.protected
            else:
                protected = None
            if protected:
                # TODO: maybe this determination belongs centralized in context?
                if dev or env.pp_host_env() in ("STAGE2", "HYPER"):
                    if ctx.ssl_client_cert_optional_in_dev:
                        sslcontext = getattr(protected, 'ssl_dev_server_context')
                    else:
                        sslcontext = getattr(protected, 'ssl_server_context')
                    server = MultiProtocolWSGIServer(
                        sock, app, spawn=10000, context=sslcontext)
                else:
                    server = SslContextWSGIServer(
                        sock, app, spawn=10000, context=protected.ssl_server_context)
            else:
                server = WSGIServer(sock, app)
            server.log = self.server_log or RotatingGeventLog()
            self.servers.append(server)
            self.socks[server] = sock
            # prevent a "blocking" call to DNS on each request
            # (NOTE: although the OS won't block, gevent will dispatch to a threadpool which is expensive)
            server.set_environ({'SERVER_NAME': gevent.socket.getfqdn(address[0])})
        for handler in stream_handlers:
            pass

    def run(self):
        if not self.prefork:
            self.start()
            try:
                while 1:
                    async.sleep(1.0)
            finally:
                self.stop()
            return
        if not ufork:
            raise RuntimeError('attempting to run pre-forked on platform without fork')
        self.arbiter = ufork.Arbiter(
            post_fork=self._post_fork, child_pre_exit=self.stop, parent_pre_stop=self.stop,
            size=self.num_workers, sleep=async.sleep, fork=gevent.fork)
        if self.daemonize:
            ctx = context.get_context()
            pidfile = os.path.join(ctx.pid_file_path,
                                   '{0}.pid'.format(ctx.appname))
            self.arbiter.spawn_daemon(pidfile=pidfile)
        else:
            self.arbiter.run()

    def _post_fork(self):
        # TODO: revisit this with newer gevent release
        hub = gevent.hub.get_hub()
        hub.loop.reinit()  # reinitializes libev
        hub._threadpool = None  # eliminate gevent's internal threadpools
        gevent.sleep(0)  # let greenlets run
        # finally, eliminate our threadpools
        context.get_context().thread_locals = threading.local()
        if self.post_fork:
            self.post_fork()
        self.start()

    def start(self):
        errs = {}
        for server in self.servers:
            try:
                server.start()
            except Exception as e:
                errs[server] = e
        if errs:
            raise RuntimeError(errs)

    def stop(self, timeout=30.0):
        async.join([async.spawn(server.stop, timeout) for server in self.servers], raise_exc=True)

def _make_server_sock(address):
    sock = gevent.socket.socket()
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(address)
    sock.listen(128)  # Note: 128 is a "hint" to the OS than a strict rule about backlog size
    return sock

def my_close(self):
    context.get_context().stats['fake_close'].add(1.0)

class SslContextWSGIServer(WSGIServer):
    def wrap_socket_and_handle(self, client_socket, address):
        if not self.ssl_args:
            raise ValueError('https server requires server-side'
                             ' SSL certificate (protected)')
        protocol = _socket_protocol(client_socket)
        if protocol == "ssl":
            ssl_socket = async.wrap_socket_context(client_socket, **self.ssl_args)
            return self.handle(ssl_socket, address)
        elif protocol == "http":
            client_socket.send(_NO_SSL_HTTP_RESPONSE)
            async.killsock(client_socket)
        else:
            context.get_context().counts["server.pings"] += 1

class MultiProtocolWSGIServer(SslContextWSGIServer):
    def wrap_socket_and_handle(self, client_socket, address):
        protocol = _socket_protocol(client_socket)
        if protocol == "http":
            return self.handle(client_socket, address)
        elif protocol == "ssl":
            ssl_socket = async.wrap_socket_context(client_socket, **self.ssl_args)
            return self.handle(ssl_socket, address)
        else:
            context.get_context().counts["server.pings"] += 1

#http://en.wikipedia.org/wiki/Hypertext_Transfer_Protocol#Request_methods
_HTTP_METHODS = ('GET', 'HEAD', 'POST', 'PUT', 'DELETE', 'TRACE', 'OPTION',
                 'CONNECT', 'PATCH')

_SSL_RECORD_TYPES = {
    '\x14': 'CHANGE_CIPHER_SPEC', '\x15': 'ALERT', '\x16': 'HANDSHAKE',
    '\x17': 'APPLICATION_DATA',
}

_SSL_VERSIONS = {
    '\x03\x00': 'SSL3', '\x03\x01': 'TLS1', '\x03\x02': 'TLS1.1', '\x03\x03': 'TLS1.2'
}

_NO_SSL_HTTP_RESPONSE = "\r\n".join([
    'HTTP/1.1 401 UNAUTHORIZED',
    'Content-Type: text/plain',
    'Content-Length: ' + str(len('SSL REQUIRED')),
    '',
    'SSL REQUIRED'
    ])  # this could probably be improved, but better than just closing the socket

def _socket_protocol(client_socket):
    first_bytes = client_socket.recv(10, socket.MSG_PEEK)
    if any([method in first_bytes for method in _HTTP_METHODS]):
        return "http"
    elif not first_bytes:
        # TODO: empty requests are sometimes sent ([socket] keepalives?)
        return
    elif first_bytes[0] in _SSL_RECORD_TYPES and first_bytes[1:3] in _SSL_VERSIONS:
        return "ssl"
    else:
        raise IOError("unrecognized client protocol: {0!r}".format(first_bytes))

class RotatingGeventLog(object):
    def __init__(self, size=10000):
        self.msgs = collections.deque()
        self.size = size

    def write(self, msg):
        self.msgs.appendleft(msg)
        if len(self.msgs) > self.size:
            self.msgs.pop()
<commit_msg>added missing import<commit_after>'''
Defines a ServerGroup class for running and managing servers, as well as a set of
base server types.

The following information is needed to start a server:
1- handler / wsgi app
2- address
3- ssl preference

'''
import collections
import socket
import os
import threading

if hasattr(os, "fork"):
    import ufork
else:
    ufork = None
from gevent.pywsgi import WSGIServer
import gevent.socket

import async
from protected import Protected
import context
import env

class ServerGroup(object):
    def __init__(self, wsgi_apps=(), stream_handlers=(), prefork=False, daemonize=False, dev=False, **kw):
        self.prefork = prefork
        self.daemonize = daemonize
        self.post_fork = kw.get('post_fork')  # callback to be executed post fork
        self.server_log = kw.get('gevent_log')
        self.wsgi_apps = wsgi_apps
        self.stream_handlers = stream_handlers
        ctx = context.get_context()
        self.num_workers = ctx.num_workers
        self.servers = []
        self.socks = {}
        for app in wsgi_apps:
            app, address, ssl = app
            sock = _make_server_sock(address)
            if isinstance(ssl, Protected):
                protected = ssl
            elif ssl:
                protected = ctx.protected
            else:
                protected = None
            if protected:
                # TODO: maybe this determination belongs centralized in context?
                if dev or env.pp_host_env() in ("STAGE2", "HYPER"):
                    if ctx.ssl_client_cert_optional_in_dev:
                        sslcontext = getattr(protected, 'ssl_dev_server_context')
                    else:
                        sslcontext = getattr(protected, 'ssl_server_context')
                    server = MultiProtocolWSGIServer(
                        sock, app, spawn=10000, context=sslcontext)
                else:
                    server = SslContextWSGIServer(
                        sock, app, spawn=10000, context=protected.ssl_server_context)
            else:
                server = WSGIServer(sock, app)
            server.log = self.server_log or RotatingGeventLog()
            self.servers.append(server)
            self.socks[server] = sock
            # prevent a "blocking" call to DNS on each request
            # (NOTE: although the OS won't block, gevent will dispatch to a threadpool which is expensive)
            server.set_environ({'SERVER_NAME': gevent.socket.getfqdn(address[0])})
        for handler in stream_handlers:
            pass

    def run(self):
        if not self.prefork:
            self.start()
            try:
                while 1:
                    async.sleep(1.0)
            finally:
                self.stop()
            return
        if not ufork:
            raise RuntimeError('attempting to run pre-forked on platform without fork')
        self.arbiter = ufork.Arbiter(
            post_fork=self._post_fork, child_pre_exit=self.stop, parent_pre_stop=self.stop,
            size=self.num_workers, sleep=async.sleep, fork=gevent.fork)
        if self.daemonize:
            ctx = context.get_context()
            pidfile = os.path.join(ctx.pid_file_path,
                                   '{0}.pid'.format(ctx.appname))
            self.arbiter.spawn_daemon(pidfile=pidfile)
        else:
            self.arbiter.run()

    def _post_fork(self):
        # TODO: revisit this with newer gevent release
        hub = gevent.hub.get_hub()
        hub.loop.reinit()  # reinitializes libev
        hub._threadpool = None  # eliminate gevent's internal threadpools
        gevent.sleep(0)  # let greenlets run
        # finally, eliminate our threadpools
        context.get_context().thread_locals = threading.local()
        if self.post_fork:
            self.post_fork()
        self.start()

    def start(self):
        errs = {}
        for server in self.servers:
            try:
                server.start()
            except Exception as e:
                errs[server] = e
        if errs:
            raise RuntimeError(errs)

    def stop(self, timeout=30.0):
        async.join([async.spawn(server.stop, timeout) for server in self.servers], raise_exc=True)

def _make_server_sock(address):
    sock = gevent.socket.socket()
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(address)
    sock.listen(128)  # Note: 128 is a "hint" to the OS than a strict rule about backlog size
    return sock

def my_close(self):
    context.get_context().stats['fake_close'].add(1.0)

class SslContextWSGIServer(WSGIServer):
    def wrap_socket_and_handle(self, client_socket, address):
        if not self.ssl_args:
            raise ValueError('https server requires server-side'
                             ' SSL certificate (protected)')
        protocol = _socket_protocol(client_socket)
        if protocol == "ssl":
            ssl_socket = async.wrap_socket_context(client_socket, **self.ssl_args)
            return self.handle(ssl_socket, address)
        elif protocol == "http":
            client_socket.send(_NO_SSL_HTTP_RESPONSE)
            async.killsock(client_socket)
        else:
            context.get_context().counts["server.pings"] += 1

class MultiProtocolWSGIServer(SslContextWSGIServer):
    def wrap_socket_and_handle(self, client_socket, address):
        protocol = _socket_protocol(client_socket)
        if protocol == "http":
            return self.handle(client_socket, address)
        elif protocol == "ssl":
            ssl_socket = async.wrap_socket_context(client_socket, **self.ssl_args)
            return self.handle(ssl_socket, address)
        else:
            context.get_context().counts["server.pings"] += 1

#http://en.wikipedia.org/wiki/Hypertext_Transfer_Protocol#Request_methods
_HTTP_METHODS = ('GET', 'HEAD', 'POST', 'PUT', 'DELETE', 'TRACE', 'OPTION',
                 'CONNECT', 'PATCH')

_SSL_RECORD_TYPES = {
    '\x14': 'CHANGE_CIPHER_SPEC', '\x15': 'ALERT', '\x16': 'HANDSHAKE',
    '\x17': 'APPLICATION_DATA',
}

_SSL_VERSIONS = {
    '\x03\x00': 'SSL3', '\x03\x01': 'TLS1', '\x03\x02': 'TLS1.1', '\x03\x03': 'TLS1.2'
}

_NO_SSL_HTTP_RESPONSE = "\r\n".join([
    'HTTP/1.1 401 UNAUTHORIZED',
    'Content-Type: text/plain',
    'Content-Length: ' + str(len('SSL REQUIRED')),
    '',
    'SSL REQUIRED'
    ])  # this could probably be improved, but better than just closing the socket

def _socket_protocol(client_socket):
    first_bytes = client_socket.recv(10, socket.MSG_PEEK)
    if any([method in first_bytes for method in _HTTP_METHODS]):
        return "http"
    elif not first_bytes:
        # TODO: empty requests are sometimes sent ([socket] keepalives?)
        return
    elif first_bytes[0] in _SSL_RECORD_TYPES and first_bytes[1:3] in _SSL_VERSIONS:
        return "ssl"
    else:
        raise IOError("unrecognized client protocol: {0!r}".format(first_bytes))

class RotatingGeventLog(object):
    def __init__(self, size=10000):
        self.msgs = collections.deque()
        self.size = size

    def write(self, msg):
        self.msgs.appendleft(msg)
        if len(self.msgs) > self.size:
            self.msgs.pop()
<|endoftext|>


<commit_before># Copyright 2016 Google Inc. All Rights Reserved.
#
# 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.

"""Tasks related to Python gRPC code generation"""

import io
import os
import re
import tempfile
import time

from ruamel import yaml

from artman.utils import protoc_utils
from artman.tasks import task_base

class PythonChangePackageTask(task_base.TaskBase):
    """Copies source protos to a package that meets Python convention"""
    default_provides = ('final_src_proto_path',
                        'final_import_proto_path')

    _IDENTIFIER = '[A-Za-z_][A-Za-z_0-9]*'

    _BASE_PROTO_REGEX = (
        '(?P<prefix>{prefix})' +
        '(?P<package>' + _IDENTIFIER +
        '({separator}' + _IDENTIFIER + ')*{package_suffix})'
        '(?P<suffix>{suffix})')

    # E.g., `package google.foo.bar`
    _PACKAGE_REGEX = re.compile(_BASE_PROTO_REGEX.format(
        prefix='^package ',
        separator='\\.',
        package_suffix='',
        suffix=''))

    # E.g., `import "google/foo/bar";`
    _IMPORT_REGEX = re.compile(_BASE_PROTO_REGEX.format(
        prefix='^import "',
        separator='/',
        package_suffix='\\.proto',
        suffix='";'))

    # TODO (geigerj): add regex for documentation link updates?

    def execute(self, src_proto_path, import_proto_path, common_protos_yaml,
                organization_name):
        self._organization_name = organization_name

        with io.open(common_protos_yaml) as file_:
            common_protos_data = yaml.load(file_, Loader=yaml.Loader)

        # Treat google.protobuf, google.iam as a common proto package, even
        # though they are not included in the common-protos we generate.
        #
        # TODO (geigerj): remove 'google.iam' when it is included in the common
        # protos package.
        common_protos = ['google.protobuf', 'google.iam']
        for package in common_protos_data['packages']:
            common_protos.append('google.' + package['name'].replace('/', '.'))

        tmpdir = os.path.join(
            tempfile.gettempdir(), 'artman-python', str(int(time.time())))
        new_proto_dir = os.path.join(tmpdir, 'proto')
        new_src_path = set()
        new_import_path = [new_proto_dir]

        self._copy_and_transform_directories(
            src_proto_path, new_proto_dir, common_protos, paths=new_src_path)
        self._copy_and_transform_directories(
            import_proto_path, new_proto_dir, common_protos)

        # Update src_proto_path, import_proto_path
        return list(new_src_path), new_import_path

    def _extract_base_dirs(self, proto_file):
        """Return the proto file path derived from the package name."""

        with io.open(proto_file, 'rt', encoding='UTF-8') as proto:
            for line in proto:
                pkg = self._PACKAGE_REGEX.match(line)
                if pkg:
                    pkg = pkg.group('package')
                    return os.path.sep.join(pkg.split('.'))
            return ''

    def _transform(self, pkg, sep, common_protos):
        """Transform to the appropriate proto package layout.

        Works with arbitrary separator (e.g., '/' for import statements,
        '.' for proto package statements, os.path.sep for filenames)
        """
        if sep != '.' and pkg.endswith('.proto'):
            dotted = pkg[:-6].replace(sep, '.')
            suffix = '.proto'
        else:
            dotted = pkg.replace(sep, '.')
            suffix = ''

        # Sanity check: Do not transform common protos.
        for common_pkg in common_protos:
            if dotted.startswith(common_pkg):
                return pkg

        # Special case: If the organization name is "google-cloud", then we
        # have to ensure that "cloud" exists in the path. The protos
        # themselves may not follow this.
        if 'cloud' not in dotted and self._organization_name == 'google-cloud':
            dotted = dotted.replace('google.', 'google.cloud.', 1)

        # Transform into the ideal proto path.
        # What essentially should happen here is that "{api}.{vN}" should
        # change to "{api}_{vN}".
        dotted = re.sub(r'\.v([\da-z_]*)([\d]+)\b', r'_v\1\2.proto', dotted)

        # Edge case: Some internal customers use "vNalpha".
        # Rather than make the regular expression more complicated, catch
        # this as a one-off.
        if re.search(r'\.v[\d]+alpha\b', dotted):
            dotted = re.sub(r'\.v([\d]+)alpha\b', r'_v\1alpha.proto', dotted)

        # Done; return with the appropriate separator.
        return dotted.replace('.', sep) + suffix

    def _copy_proto(self, src, dest, common_protos):
        """Copies a proto while fixing its imports"""
        with io.open(src, 'r', encoding='UTF-8') as src_lines:
            with io.open(dest, 'w+', encoding='UTF-8') as dest_file:
                for line in src_lines:
                    import_ = self._IMPORT_REGEX.match(line)
                    if import_:
                        dest_file.write('import "{}";\n'.format(
                            self._transform(
                                import_.group('package'), '/', common_protos)))
                    else:
                        dest_file.write(line)

    def _copy_and_transform_directories(
            self, src_directories, destination_directory, common_protos,
            paths=None):
        for path in src_directories:
            protos = list(protoc_utils.find_protos([path], []))
            for proto in protos:
                src_base_dirs = self._extract_base_dirs(proto)
                sub_new_src = os.path.join(
                    destination_directory,
                    self._transform(
                        src_base_dirs, os.path.sep, common_protos))
                if paths is not None:
                    paths.add(sub_new_src)

                dest = os.path.join(sub_new_src, os.path.basename(proto))
                if not os.path.exists(dest):
                    self.exec_command(['mkdir', '-p', sub_new_src])
                self._copy_proto(
                    proto, os.path.join(sub_new_src, dest), common_protos)

class PythonMoveProtosTask(task_base.TaskBase):
    default_provides = {'grpc_code_dir'}

    def execute(self, grpc_code_dir, gapic_code_dir):
        """Move the protos into the GAPIC structure.

        This copies the ``x/y/z/proto/`` directory over to be a sibling
        of ``x/y/z/gapic/`` in the GAPIC code directory. In the event of
        an inconsistency on the prefix, the GAPIC wins.

        Args:
            grpc_code_dir (str): The location where the GRPC code was
                generated.
            gapic_code_dir (str): The location where the GAPIC code was
                generated.
        """
        # Determine the appropriate source and target directory.
        # We can get this by drilling in to the GAPIC artifact until we get to
        # a "gapic" directory.
        src = self._get_subdir_path(grpc_code_dir, 'proto')
        target = self._get_subdir_path(
            os.path.join(gapic_code_dir, 'google'),
            'gapic',
        )

        # Move the contents into the GAPIC directory.
        self.exec_command(['mv', os.path.join(src, 'proto'), target])

        # Create an __init__.py file in the proto directory.
        # This is necessary for Python 2.7 compatibility.
        self.exec_command([
            'touch', os.path.join(target, 'proto', '__init__.py'),
        ])

        # Remove the grpc directory.
        self.exec_command(['rm', '-rf', grpc_code_dir])

        # Clear out the grpc_code_dir, so future tasks perceive it as
        # not being a thing anymore.
        return {'grpc_code_dir': None}

    def _get_subdir_path(self, haystack, needle):
        """Return the subpath which contains the ``needle`` directory.

        Args:
            haystack (str): The top-level directory in which the subdirectory
                should appear.
            needle (str): The directory being sought.

        Returns:
            str: The path, relative to ``haystack``, where the subdirectory
                was found.

        Raises:
            RuntimeError: If the subdirectory is not found.
        """
        for path, dirs, files in os.walk(haystack):
            if needle in dirs:
                return path
        raise RuntimeError('Path %s not found in %s.' % (needle, haystack))
<commit_msg>Allow "v1beta" as a version name. (#359)<commit_after># Copyright 2016 Google Inc. All Rights Reserved.
#
# 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.

"""Tasks related to Python gRPC code generation"""

import io
import os
import re
import tempfile
import time

from ruamel import yaml

from artman.utils import protoc_utils
from artman.tasks import task_base

class PythonChangePackageTask(task_base.TaskBase):
    """Copies source protos to a package that meets Python convention"""
    default_provides = ('final_src_proto_path',
                        'final_import_proto_path')

    _IDENTIFIER = '[A-Za-z_][A-Za-z_0-9]*'

    _BASE_PROTO_REGEX = (
        '(?P<prefix>{prefix})' +
        '(?P<package>' + _IDENTIFIER +
        '({separator}' + _IDENTIFIER + ')*{package_suffix})'
        '(?P<suffix>{suffix})')

    # E.g., `package google.foo.bar`
    _PACKAGE_REGEX = re.compile(_BASE_PROTO_REGEX.format(
        prefix='^package ',
        separator='\\.',
        package_suffix='',
        suffix=''))

    # E.g., `import "google/foo/bar";`
    _IMPORT_REGEX = re.compile(_BASE_PROTO_REGEX.format(
        prefix='^import "',
        separator='/',
        package_suffix='\\.proto',
        suffix='";'))

    # TODO (geigerj): add regex for documentation link updates?

    def execute(self, src_proto_path, import_proto_path, common_protos_yaml,
                organization_name):
        self._organization_name = organization_name

        with io.open(common_protos_yaml) as file_:
            common_protos_data = yaml.load(file_, Loader=yaml.Loader)

        # Treat google.protobuf, google.iam as a common proto package, even
        # though they are not included in the common-protos we generate.
        #
        # TODO (geigerj): remove 'google.iam' when it is included in the common
        # protos package.
        common_protos = ['google.protobuf', 'google.iam']
        for package in common_protos_data['packages']:
            common_protos.append('google.' + package['name'].replace('/', '.'))

        tmpdir = os.path.join(
            tempfile.gettempdir(), 'artman-python', str(int(time.time())))
        new_proto_dir = os.path.join(tmpdir, 'proto')
        new_src_path = set()
        new_import_path = [new_proto_dir]

        self._copy_and_transform_directories(
            src_proto_path, new_proto_dir, common_protos, paths=new_src_path)
        self._copy_and_transform_directories(
            import_proto_path, new_proto_dir, common_protos)

        # Update src_proto_path, import_proto_path
        return list(new_src_path), new_import_path

    def _extract_base_dirs(self, proto_file):
        """Return the proto file path derived from the package name."""

        with io.open(proto_file, 'rt', encoding='UTF-8') as proto:
            for line in proto:
                pkg = self._PACKAGE_REGEX.match(line)
                if pkg:
                    pkg = pkg.group('package')
                    return os.path.sep.join(pkg.split('.'))
            return ''

    def _transform(self, pkg, sep, common_protos):
        """Transform to the appropriate proto package layout.

        Works with arbitrary separator (e.g., '/' for import statements,
        '.' for proto package statements, os.path.sep for filenames)
        """
        if sep != '.' and pkg.endswith('.proto'):
            dotted = pkg[:-6].replace(sep, '.')
            suffix = '.proto'
        else:
            dotted = pkg.replace(sep, '.')
            suffix = ''

        # Sanity check: Do not transform common protos.
        for common_pkg in common_protos:
            if dotted.startswith(common_pkg):
                return pkg

        # Special case: If the organization name is "google-cloud", then we
        # have to ensure that "cloud" exists in the path. The protos
        # themselves may not follow this.
        if 'cloud' not in dotted and self._organization_name == 'google-cloud':
            dotted = dotted.replace('google.', 'google.cloud.', 1)

        # Transform into the ideal proto path.
        # What essentially should happen here is that "{api}.{vN}" should
        # change to "{api}_{vN}".
        dotted = re.sub(r'\.v([\da-z_]*)([\d]+)\b', r'_v\1\2.proto', dotted)

        # Edge case: Some internal customers use "vNalpha" and "vNbeta".
        # Rather than make the regular expression more complicated, catch
        # this as a one-off.
        if re.search(r'\.v[\d]+alpha\b', dotted):
            dotted = re.sub(r'\.v([\d]+)alpha\b', r'_v\1alpha.proto', dotted)
        if re.search(r'\.v[\d]+beta\b', dotted):
            dotted = re.sub(r'\.v([\d]+)beta\b', r'_v\1beta.proto', dotted)

        # Done; return with the appropriate separator.
        return dotted.replace('.', sep) + suffix

    def _copy_proto(self, src, dest, common_protos):
        """Copies a proto while fixing its imports"""
        with io.open(src, 'r', encoding='UTF-8') as src_lines:
            with io.open(dest, 'w+', encoding='UTF-8') as dest_file:
                for line in src_lines:
                    import_ = self._IMPORT_REGEX.match(line)
                    if import_:
                        dest_file.write('import "{}";\n'.format(
                            self._transform(
                                import_.group('package'), '/', common_protos)))
                    else:
                        dest_file.write(line)

    def _copy_and_transform_directories(
            self, src_directories, destination_directory, common_protos,
            paths=None):
        for path in src_directories:
            protos = list(protoc_utils.find_protos([path], []))
            for proto in protos:
                src_base_dirs = self._extract_base_dirs(proto)
                sub_new_src = os.path.join(
                    destination_directory,
                    self._transform(
                        src_base_dirs, os.path.sep, common_protos))
                if paths is not None:
                    paths.add(sub_new_src)

                dest = os.path.join(sub_new_src, os.path.basename(proto))
                if not os.path.exists(dest):
                    self.exec_command(['mkdir', '-p', sub_new_src])
                self._copy_proto(
                    proto, os.path.join(sub_new_src, dest), common_protos)

class PythonMoveProtosTask(task_base.TaskBase):
    default_provides = {'grpc_code_dir'}

    def execute(self, grpc_code_dir, gapic_code_dir):
        """Move the protos into the GAPIC structure.

        This copies the ``x/y/z/proto/`` directory over to be a sibling
        of ``x/y/z/gapic/`` in the GAPIC code directory. In the event of
        an inconsistency on the prefix, the GAPIC wins.

        Args:
            grpc_code_dir (str): The location where the GRPC code was
                generated.
            gapic_code_dir (str): The location where the GAPIC code was
                generated.
        """
        # Determine the appropriate source and target directory.
        # We can get this by drilling in to the GAPIC artifact until we get to
        # a "gapic" directory.
        src = self._get_subdir_path(grpc_code_dir, 'proto')
        target = self._get_subdir_path(
            os.path.join(gapic_code_dir, 'google'),
            'gapic',
        )

        # Move the contents into the GAPIC directory.
        self.exec_command(['mv', os.path.join(src, 'proto'), target])

        # Create an __init__.py file in the proto directory.
        # This is necessary for Python 2.7 compatibility.
        self.exec_command([
            'touch', os.path.join(target, 'proto', '__init__.py'),
        ])

        # Remove the grpc directory.
        self.exec_command(['rm', '-rf', grpc_code_dir])

        # Clear out the grpc_code_dir, so future tasks perceive it as
        # not being a thing anymore.
        return {'grpc_code_dir': None}

    def _get_subdir_path(self, haystack, needle):
        """Return the subpath which contains the ``needle`` directory.

        Args:
            haystack (str): The top-level directory in which the subdirectory
                should appear.
            needle (str): The directory being sought.

        Returns:
            str: The path, relative to ``haystack``, where the subdirectory
                was found.

        Raises:
            RuntimeError: If the subdirectory is not found.
        """
        for path, dirs, files in os.walk(haystack):
            if needle in dirs:
                return path
        raise RuntimeError('Path %s not found in %s.' % (needle, haystack))
<|endoftext|>


<commit_before>"""
module for making rig controls
"""

import maya.cmds as pm
import pymel.core as pm
import mayaLib.pipelineLib.utility.nameCheck as nc
from mayaLib.rigLib.utils import ctrlShape

class Control():
    """
    class for building rig control
    """

    def __init__(
            self,
            prefix='new',
            scale=1.0,
            translateTo='',
            rotateTo='',
            parent='',
            shape='circle',
            lockChannels=['s', 'v'],
            doOffset=True,
            doModify=False
    ):

        """
        :param prefix: str, prefix to name new objects
        :param scale: float, scale value for size of control shapes
        :param translateTo: str, reference object for control position
        :param rotateTo: str, reference object for control orientation
        :param parent: str, object to be parent of new control
        :param shape: str, control shape type (normal direction)
        :param lockChannels: list( str ), list of channels on control to be locked and non-keyable
        :return: None
        """

        # name handle
        if '*' in prefix:
            prefix = nc.nameCheck(prefix + '_CTRL').split('_')[0]

        ctrlObject = None
        circleNormal = [1, 0, 0]

        if shape in ['circle', 'circleX']:
            circleNormal = [1, 0, 0]

        elif shape == 'circleY':
            circleNormal = [0, 1, 0]

        elif shape == 'circleZ':
            circleNormal = [0, 0, 1]

        elif shape == 'sphere':
            ctrlObject = pm.sphere(n=prefix + '_CTRL', ch=False, normal=[1, 0, 0], radius=scale)[0]

        elif shape == 'move':
            ctrlObject = ctrlShape.moveCtrlShape(name=prefix + '_CTRL', scale=scale)

        if not ctrlObject:
            ctrlObject = pm.circle(n=prefix + '_CTRL', ch=False, normal=circleNormal, radius=scale)[0]

        if doModify:
            ctrlModify = pm.group(n=prefix + 'Modify_GRP', em=1)
            pm.parent(ctrlObject, ctrlModify)

        if doOffset:
            ctrlOffset = pm.group(n=prefix + 'Offset_GRP', em=1)
            if doModify:
                pm.parent(ctrlModify, ctrlOffset)
            else:
                pm.parent(ctrlObject, ctrlOffset)

        # color control
        ctrlShapes = ctrlObject.getShape()
        ctrlShapes.ove.set(1)

        if prefix.startswith('l_'):
            ctrlShapes.ovc.set(6)

        elif prefix.startswith('r_'):
            ctrlShapes.ovc.set(13)

        else:
            ctrlShapes.ovc.set(22)

        # translate control
        if translateTo != None and translateTo != '':
            if pm.objExists(translateTo):
                pm.delete(pm.pointConstraint(translateTo, ctrlOffset))

        # rotate control
        if rotateTo != None and rotateTo != '':
            if pm.objExists(rotateTo):
                pm.delete(pm.orientConstraint(rotateTo, ctrlOffset))

        # lock control channels

        singleAttributeLockList = []

        for lockChannel in lockChannels:

            if lockChannel in ['t', 'r', 's']:

                for axis in ['x', 'y', 'z']:
                    at = lockChannel + axis
                    singleAttributeLockList.append(at)

            else:

                singleAttributeLockList.append(lockChannel)

        for at in singleAttributeLockList:
            pm.setAttr(ctrlObject + '.' + at, l=1, k=0)

        # add public members

        self.C = ctrlObject
        self.Modify = None
        self.Off = None

        if doOffset:
            self.Off = ctrlOffset
        if doModify:
            self.Modify = ctrlModify

        # parent control
        if parent != None and parent != '':
            if pm.objExists(parent):
                pm.parent(self.getTop(), parent)

    def getControl(self):
        """
        Return Control
        :return:
        """
        return self.C

    def getOffsetGrp(self):
        """
        Return Offset Grp if exist or None
        :return:
        """
        if self.Off:
            return self.Off
        return None

    def getModifyGrp(self):
        """
        Return Modify Grp if exist or None
        :return:
        """
        if self.Modify:
            return self.Modify
        return None

    def getTop(self):
        """
        Return control's top Grp or Control
        :return:
        """
        if self.Off:
            return self.Off
        elif self.Modify:
            return self.Modify
        else:
            return self.C

<commit_msg>control class update<commit_after>"""
module for making rig controls
"""

import pymel.core as pm
import mayaLib.pipelineLib.utility.nameCheck as nc
from mayaLib.rigLib.utils import ctrlShape
from mayaLib.rigLib.utils import common

class Control():
    """
    class for building rig control
    """

    def __init__(
            self,
            prefix='new',
            scale=1.0,
            translateTo='',
            rotateTo='',
            parent='',
            shape='circle',
            lockChannels=['s', 'v'],
            doOffset=True,
            doModify=False
    ):

        """
        :param prefix: str, prefix to name new objects
        :param scale: float, scale value for size of control shapes
        :param translateTo: str, reference object for control position
        :param rotateTo: str, reference object for control orientation
        :param parent: str, object to be parent of new control
        :param shape: str, control shape type (normal direction)
        :param lockChannels: list( str ), list of channels on control to be locked and non-keyable
        :return: None
        """

        # name handle
        if '*' in prefix:
            prefix = nc.nameCheck(prefix + '_CTRL').split('_')[0]

        ctrlObject = None
        circleNormal = [1, 0, 0]

        if shape in ['circle', 'circleX']:
            circleNormal = [1, 0, 0]

        elif shape == 'circleY':
            circleNormal = [0, 1, 0]

        elif shape == 'circleZ':
            circleNormal = [0, 0, 1]

        elif shape == 'sphere':
            ctrlObject = pm.sphere(n=prefix + '_CTRL', ch=False, normal=[1, 0, 0], radius=scale)[0]

        elif shape == 'move':
            ctrlObject = ctrlShape.moveCtrlShape(name=prefix + '_CTRL', scale=scale)

        elif shape == 'spine':
            ctrlObject = ctrlShape.trapeziumCtrlShape(name=prefix + '_CTRL', scale=scale)
            ctrlObject.translateY.set(3 * scale)
            common.freezeTranform(ctrlObject)

        if not ctrlObject:
            ctrlObject = pm.circle(n=prefix + '_CTRL', ch=False, normal=circleNormal, radius=scale)[0]

        if doModify:
            ctrlModify = pm.group(n=prefix + 'Modify_GRP', em=1)
            pm.parent(ctrlObject, ctrlModify)

        if doOffset:
            ctrlOffset = pm.group(n=prefix + 'Offset_GRP', em=1)
            if doModify:
                pm.parent(ctrlModify, ctrlOffset)
            else:
                pm.parent(ctrlObject, ctrlOffset)

        # color control
        ctrlShapes = ctrlObject.getShape()
        ctrlShapes.ove.set(1)

        if prefix.startswith('l_'):
            ctrlShapes.ovc.set(6)

        elif prefix.startswith('r_'):
            ctrlShapes.ovc.set(13)

        else:
            ctrlShapes.ovc.set(22)

        # translate control
        if translateTo != None and translateTo != '':
            if pm.objExists(translateTo):
                pm.delete(pm.pointConstraint(translateTo, ctrlOffset))

        # rotate control
        if rotateTo != None and rotateTo != '':
            if pm.objExists(rotateTo):
                pm.delete(pm.orientConstraint(rotateTo, ctrlOffset))

        # lock control channels

        singleAttributeLockList = []

        for lockChannel in lockChannels:

            if lockChannel in ['t', 'r', 's']:

                for axis in ['x', 'y', 'z']:
                    at = lockChannel + axis
                    singleAttributeLockList.append(at)

            else:

                singleAttributeLockList.append(lockChannel)

        for at in singleAttributeLockList:
            pm.setAttr(ctrlObject + '.' + at, l=1, k=0)

        # add public members

        self.C = ctrlObject
        self.Modify = None
        self.Off = None

        if doOffset:
            self.Off = ctrlOffset
        if doModify:
            self.Modify = ctrlModify

        # parent control
        if parent != None and parent != '':
            if pm.objExists(parent):
                pm.parent(self.getTop(), parent)

    def getControl(self):
        """
        Return Control
        :return:
        """
        return self.C

    def getOffsetGrp(self):
        """
        Return Offset Grp if exist or None
        :return:
        """
        if self.Off:
            return self.Off
        return None

    def getModifyGrp(self):
        """
        Return Modify Grp if exist or None
        :return:
        """
        if self.Modify:
            return self.Modify
        return None

    def getTop(self):
        """
        Return control's top Grp or Control
        :return:
        """
        if self.Off:
            return self.Off
        elif self.Modify:
            return self.Modify
        else:
            return self.C

<|endoftext|>


<commit_before>from __future__ import division, print_function
import argparse
from glob import glob
import os

import numpy as np

from .featurize import get_embeddings
from .reader import VERSIONS
from .stats import load_stats, save_stats
from .sort import sort_by_region

def main():
    parser = argparse.ArgumentParser(
        description="Reads American Community Survey Public Use Microdata "
                    "Sample files and featurizes them, as in Flaxman, Wang, "
                    "and Smola (KDD 2015). Needs a preprocessing pass over "
                    "the data to sort it into files by region and to collect "
                    "statistics in order to do one-hot encoding and "
                    "standardization. Currently supports only the 2006-10 "
                    "file used in the original paper.",
    )
    subparsers = parser.add_subparsers()

    ############################################################################
    sort = subparsers.add_parser(
        'sort', help="Sort the data by region and collect statistics about it.")
    sort.set_defaults(func=do_sort)

    io = sort.add_argument_group('Input/output options')
    g = io.add_mutually_exclusive_group(required=True)
    g.add_argument('--zipfile', '-z', help="The original ACS PUMS zip file.")
    g.add_argument('--csv-files', '-c', nargs='+',
                   help="CSV files in ACS PUMS format.")

    io.add_argument('out_dir', help='Directory for the sorted features.')

    io.add_argument('--chunksize', type=int, default=10**5, metavar='LINES',
                      help="How much of a CSV file to read at a time; default "
                           "%(default)s.")
    io.add_argument('--stats-only', action='store_true', default=False,
                    help="Only compute the stats, don't do the sorting.")

    fmt = sort.add_argument_group('Format options')
    g = fmt.add_mutually_exclusive_group()
    g.add_argument('--voters-only', action='store_true', default=True,
                   help="Only include citizens who are at least 18 years old "
                        "(default).")
    g.add_argument('--all-people', action='store_false', dest='voters_only',
                   help="Include all records from the files.")
    fmt.add_argument('--version', choices=VERSIONS, default='2006-10',
                      help="The format of the ACS PUMS files in use; default "
                           "%(default)s.")

    ############################################################################
    featurize = subparsers.add_parser(
        'featurize', help="Emit features for a given region.")
    featurize.set_defaults(func=do_featurize)

    io = featurize.add_argument_group('Input/output options')
    io.add_argument('dir', help="The directory where `pummel sort` put stuff.")
    io.add_argument('outfile', nargs='?',
                    help='Where to put embeddings; default DIR/embeddings.npz.')
    io.add_argument('--chunksize', type=int, default=2**13, metavar='LINES',
                      help="How much of a region to process at a time; default "
                           "%(default)s.")

    emb = featurize.add_argument_group('Embedding options')
    emb.add_argument('--skip-rbf', action='store_true', default=False,
                     help="Skip getting the RBF kernel embedding and only get "
                          "the linear one (much, much faster).")
    emb.add_argument('--n-freqs', type=int, default=2048,
                     help='Number of random frequencies to use (half the '
                          'embedding dimension; default %(default)s).')
    emb.add_argument('--bandwidth', type=float,
                     help='Gaussian kernel bandwidth. Default: choose the '
                          'median distance among the random sample saved in '
                          'the stats file.')
    emb.add_argument('--skip-feats', nargs='+', metavar='FEAT_NAME',
                     help="Don't include some features in the embedding.")

    ############################################################################
    export = subparsers.add_parser(
        'export', help="Export features in embeddings.npz as CSV files.")
    io2 = export.add_argument_group('Input/output options')
    io2.add_argument('dir', help="The directory where `pummel featurize` put stuff.")
    io2.add_argument('infile', nargs='?',
                    help='Location of embeddings created by `pummel embeddings`; default DIR/embeddings.npz.')
    export.set_defaults(func=do_export)

    ############################################################################
    args = parser.parse_args()
    args.func(args, parser)

def do_sort(args, parser):
    if not os.path.isdir(args.out_dir):
        os.makedirs(args.out_dir)
    stats = sort_by_region(
        args.zipfile or args.csv_files,
        os.path.join(args.out_dir, 'feats_{}.h5'),
        voters_only=args.voters_only, stats_only=args.stats_only,
        adj_inc=True, version=args.version, chunksize=args.chunksize)
    save_stats(os.path.join(args.out_dir, 'stats.h5'), stats)

def do_export(args, parser):
    if args.infile is None:
        args.infile = os.path.join(args.dir, 'embeddings.npz')
    data = np.load(args.infile)
    df = pd.DataFrame(data['emb_lin'])
    df.set_index(data['region_names'],inplace=True)
    df.columns = data['feature_names']
    df.to_csv(os.path.join(args.dir,'embeddings_linear.csv'), index_label="region")
    print("linear embeddings saved in",os.path.join(args.dir,'embeddings_linear.csv'))

    if 'emb_rff' in data:
        df = pd.DataFrame(data['emb_rff'])
        df.set_index(data['region_names'],inplace=True)
        df.to_csv(os.path.join(args.dir,'embeddings_rff.csv'), index_label="region")
        print("rff embeddings saved in",os.path.join(args.dir,'embeddings_rff.csv'))

def do_featurize(args, parser):
    if args.outfile is None:
        args.outfile = os.path.join(args.dir, 'embeddings.npz')
    stats = load_stats(os.path.join(args.dir, 'stats.h5'))
    files = glob(os.path.join(args.dir, 'feats_*.h5'))
    region_names = [os.path.basename(f)[6:-3] for f in files]
    if args.skip_rbf:
        emb_lin, feature_names = get_embeddings(
            files, stats=stats, chunksize=args.chunksize, skip_rbf=True,
            skip_feats=args.skip_feats)
        np.savez(args.outfile, emb_lin=emb_lin,
                 feature_names=feature_names, region_names=region_names)
    else:
        emb_lin, emb_rff, freqs, bandwidth, feature_names = get_embeddings(
            files, stats=stats, n_freqs=args.n_freqs, bandwidth=args.bandwidth,
            skip_feats=args.skip_feats, chunksize=args.chunksize)
        np.savez(args.outfile,
                 emb_lin=emb_lin, emb_rff=emb_rff,
                 freqs=freqs, bandwidth=bandwidth,
                 feature_names=feature_names, region_names=region_names)
<commit_msg>a few fixes/improvements to export subcommand<commit_after>from __future__ import division, print_function
import argparse
from glob import glob
import os

import numpy as np
import pandas as pd

from .featurize import get_embeddings
from .reader import VERSIONS
from .stats import load_stats, save_stats
from .sort import sort_by_region

def main():
    parser = argparse.ArgumentParser(
        description="Reads American Community Survey Public Use Microdata "
                    "Sample files and featurizes them, as in Flaxman, Wang, "
                    "and Smola (KDD 2015). Needs a preprocessing pass over "
                    "the data to sort it into files by region and to collect "
                    "statistics in order to do one-hot encoding and "
                    "standardization. Currently supports only the 2006-10 "
                    "file used in the original paper.",
    )
    subparsers = parser.add_subparsers()

    ############################################################################
    sort = subparsers.add_parser(
        'sort', help="Sort the data by region and collect statistics about it.")
    sort.set_defaults(func=do_sort)

    io = sort.add_argument_group('Input/output options')
    g = io.add_mutually_exclusive_group(required=True)
    g.add_argument('--zipfile', '-z', help="The original ACS PUMS zip file.")
    g.add_argument('--csv-files', '-c', nargs='+',
                   help="CSV files in ACS PUMS format.")

    io.add_argument('out_dir', help='Directory for the sorted features.')

    io.add_argument('--chunksize', type=int, default=10**5, metavar='LINES',
                      help="How much of a CSV file to read at a time; default "
                           "%(default)s.")
    io.add_argument('--stats-only', action='store_true', default=False,
                    help="Only compute the stats, don't do the sorting.")

    fmt = sort.add_argument_group('Format options')
    g = fmt.add_mutually_exclusive_group()
    g.add_argument('--voters-only', action='store_true', default=True,
                   help="Only include citizens who are at least 18 years old "
                        "(default).")
    g.add_argument('--all-people', action='store_false', dest='voters_only',
                   help="Include all records from the files.")
    fmt.add_argument('--version', choices=VERSIONS, default='2006-10',
                      help="The format of the ACS PUMS files in use; default "
                           "%(default)s.")

    ############################################################################
    featurize = subparsers.add_parser(
        'featurize', help="Emit features for a given region.")
    featurize.set_defaults(func=do_featurize)

    io = featurize.add_argument_group('Input/output options')
    io.add_argument('dir', help="The directory where `pummel sort` put stuff.")
    io.add_argument('outfile', nargs='?',
                    help='Where to put embeddings; default DIR/embeddings.npz.')
    io.add_argument('--chunksize', type=int, default=2**13, metavar='LINES',
                      help="How much of a region to process at a time; default "
                           "%(default)s.")

    emb = featurize.add_argument_group('Embedding options')
    emb.add_argument('--skip-rbf', action='store_true', default=False,
                     help="Skip getting the RBF kernel embedding and only get "
                          "the linear one (much, much faster).")
    emb.add_argument('--n-freqs', type=int, default=2048,
                     help='Number of random frequencies to use (half the '
                          'embedding dimension; default %(default)s).')
    emb.add_argument('--bandwidth', type=float,
                     help='Gaussian kernel bandwidth. Default: choose the '
                          'median distance among the random sample saved in '
                          'the stats file.')
    emb.add_argument('--skip-feats', nargs='+', metavar='FEAT_NAME',
                     help="Don't include some features in the embedding.")

    ############################################################################
    export = subparsers.add_parser(
        'export', help="Export features in embeddings.npz as CSV files.")
    export.set_defaults(func=do_export)

    io = export.add_argument_group('Input/output options')
    io.add_argument('dir', help="Where to put the outputs.")
    io.add_argument('infile', nargs='?',
                    help="Location of embeddings created by `pummel feauturize`"
                         "; default DIR/embeddings.npz.")
    io.add_argument('--out-name', metavar='BASE',
                    help="Prefix for embedding output files, so that they "
                         "go e.g. in DIR/BASE_linear.csv. Default to "
                         "the basename of INFILE if it's in DIR or "
                         "otherwise 'embeddings'.")

    ############################################################################
    args = parser.parse_args()
    args.func(args, parser)

def do_sort(args, parser):
    if not os.path.isdir(args.out_dir):
        os.makedirs(args.out_dir)
    stats = sort_by_region(
        args.zipfile or args.csv_files,
        os.path.join(args.out_dir, 'feats_{}.h5'),
        voters_only=args.voters_only, stats_only=args.stats_only,
        adj_inc=True, version=args.version, chunksize=args.chunksize)
    save_stats(os.path.join(args.out_dir, 'stats.h5'), stats)

def do_featurize(args, parser):
    if args.outfile is None:
        args.outfile = os.path.join(args.dir, 'embeddings.npz')
    stats = load_stats(os.path.join(args.dir, 'stats.h5'))
    files = glob(os.path.join(args.dir, 'feats_*.h5'))
    region_names = [os.path.basename(f)[6:-3] for f in files]
    if args.skip_rbf:
        emb_lin, feature_names = get_embeddings(
            files, stats=stats, chunksize=args.chunksize, skip_rbf=True,
            skip_feats=args.skip_feats)
        np.savez(args.outfile, emb_lin=emb_lin,
                 feature_names=feature_names, region_names=region_names)
    else:
        emb_lin, emb_rff, freqs, bandwidth, feature_names = get_embeddings(
            files, stats=stats, n_freqs=args.n_freqs, bandwidth=args.bandwidth,
            skip_feats=args.skip_feats, chunksize=args.chunksize)
        np.savez(args.outfile,
                 emb_lin=emb_lin, emb_rff=emb_rff,
                 freqs=freqs, bandwidth=bandwidth,
                 feature_names=feature_names, region_names=region_names)

def do_export(args, parser):
    if args.infile is None:
        args.infile = os.path.join(args.dir, 'embeddings.npz')

    if args.out_name is None:
        rel = os.path.relpath(args.infile, args.dir)
        if '/' in rel:
            args.out_name = 'embeddings'
        else:
            args.out_name = rel[:-4] if rel.endswith('.npz') else rel
    out_pattern = os.path.join(args.dir, args.out_name + '_{}.csv')

    with np.load(args.infile) as data:
        path = out_pattern.format('linear')
        df = pd.DataFrame(data['emb_lin'])
        df.set_index(data['region_names'], inplace=True)
        df.columns = data['feature_names']
        df.to_csv(path, index_label="region")
        print("Linear embeddings saved in {}".format(path))

        if 'emb_rff' in data:
            path = out_pattern.format('rff')
            df = pd.DataFrame(data['emb_rff'])
            df.set_index(data['region_names'], inplace=True)
            df.to_csv(path, index_label="region")
            print("Fourier embeddings saved in {}".format(path))
<|endoftext|>


<commit_before>class Collection(object):
    """See 3.3.5 Emulating container types: http://docs.python.org/ref/sequence-types.html#l2h-232"""

    def __init__(self, g, uri, seq=[]):
        self.g = g
        self.uri = uri or BNode()
        for item in seq:
            self.append(item)

    def _get_container(self, index):
        """Gets the first, rest holding node at index."""
        assert isinstance(index, int)
        container = self.uri
        i = 0
        while i<index:
            i += 1
            container = g.value(container, RDF.rest)
            if container is None:
                break
        return container

    def __len__(self):
        """length of items in collection."""
        count = 0
        for item in g.items(self.uri):
            count += 1
        return count

    def __getitem__(self, key):
        """..."""
        #print "Get item %s list." % key
        c = self._get_container(key)
        if c:
            v = g.value(c, RDF.first)
            if v:
                return v
            else:
                raise KeyError, key
        else:
            raise IndexError, key

    def __setitem__(self, key, value):
        """..."""
        #print "Set item %s to %s list." % (key, value)
        c = self._get_container(key)
        if c:
            g.add((c, RDF.first, value))
        else:
            raise IndexError, key

    def __delitem__(self, key):
        """..."""
        #print "Delete item %s list." % key
        self[key] # to raise any potential key exceptions
        current = self._get_container(key)
        assert current
        next = self._get_container(key+1)
        # TODO: delete last element in list
        assert next
        first = g.value(next, RDF.first)
        rest = g.value(next, RDF.rest)

        g.set((current, RDF.first, first))
        g.set((current, RDF.rest, rest))

    def __iter__(self):
        """Iterator over items in Collections"""
        return g.items(self.uri)

    def append(self, item):
        container = self.uri
        while True:
            first = g.value(container, RDF.first)
            if first is None:
                g.add((container, RDF.first, item))
                return
            else:
                rest = g.value(container, RDF.rest)
                if rest:
                    container = rest
                else:
                    node = BNode()
                    g.add((container, RDF.rest, node))
                    container = node

if __name__=="__main__":
    from rdflib import *

    g = Graph()

    c = Collection(g, BNode())

    assert len(c)==0

    c = Collection(g, BNode(), [Literal("1"), Literal("2"), Literal("3"), Literal("4")])

    assert len(c)==4

    assert c[1]==Literal("2"), c[1]

    del c[1]

    assert list(c)==[Literal("1"), Literal("3"), Literal("4")], list(c)

    try:
        del c[500]
    except IndexError, i:
        pass

    c.append(Literal("5"))

    print list(c)

    for i in c:
        print i

    del c[3]
<commit_msg>minor cleanup<commit_after>
class Collection(object):
    """See 3.3.5 Emulating container types: http://docs.python.org/ref/sequence-types.html#l2h-232"""

    def __init__(self, graph, uri, seq=[]):
        self.graph = graph
        self.uri = uri or BNode()
        for item in seq:
            self.append(item)

    def _get_container(self, index):
        """Gets the first, rest holding node at index."""
        assert isinstance(index, int)
        graph = self.graph
        container = self.uri
        i = 0
        while i<index:
            i += 1
            container = graph.value(container, RDF.rest)
            if container is None:
                break
        return container

    def __len__(self):
        """length of items in collection."""
        count = 0
        for item in self.graph.items(self.uri):
            count += 1
        return count

    def __getitem__(self, key):
        """TODO"""
        c = self._get_container(key)
        if c:
            v = self.graph.value(c, RDF.first)
            if v:
                return v
            else:
                raise KeyError, key
        else:
            raise IndexError, key

    def __setitem__(self, key, value):
        """TODO"""
        c = self._get_container(key)
        if c:
            self.graph.add((c, RDF.first, value))
        else:
            raise IndexError, key

    def __delitem__(self, key):
        """..."""
        self[key] # to raise any potential key exceptions
        graph = self.graph
        current = self._get_container(key)
        assert current
        next = self._get_container(key+1)
        # TODO: delete last element in list
        assert next
        first = graph.value(next, RDF.first)
        rest = graph.value(next, RDF.rest)

        graph.set((current, RDF.first, first))
        graph.set((current, RDF.rest, rest))

    def __iter__(self):
        """Iterator over items in Collections"""
        return self.graph.items(self.uri)

    def append(self, item):
        container = self.uri
        graph = self.graph
        while True:
            first = graph.value(container, RDF.first)
            if first is None:
                graph.add((container, RDF.first, item))
                return
            else:
                rest = graph.value(container, RDF.rest)
                if rest:
                    container = rest
                else:
                    node = BNode()
                    graph.add((container, RDF.rest, node))
                    container = node

if __name__=="__main__":
    from rdflib import *

    g = Graph()

    c = Collection(g, BNode())

    assert len(c)==0

    c = Collection(g, BNode(), [Literal("1"), Literal("2"), Literal("3"), Literal("4")])

    assert len(c)==4

    assert c[1]==Literal("2"), c[1]

    del c[1]

    assert list(c)==[Literal("1"), Literal("3"), Literal("4")], list(c)

    try:
        del c[500]
    except IndexError, i:
        pass

    c.append(Literal("5"))

    print list(c)

    for i in c:
        print i

    del c[3]
<|endoftext|>


<commit_before>__author__ = 'Lorenzo Argentieri'

import pymel.core as pm
from maya import mel

from mayaLib.rigLib.utils import name

def createFollicle(geo, u, v, prefix):
    follicleShape = pm.createNode('follicle', n=prefix+'_FLCShape')
    geoShape = geo.getShape()
    follicle = follicleShape.getParent()
    follicle.rename(prefix + '_FLC')

    pm.connectAttr(geoShape.outMesh, follicleShape.inputMesh)
    pm.connectAttr(geoShape.worldMatrix, follicleShape.inputWorldMatrix)

    pm.connectAttr(follicleShape.outRotate, follicle.rotate)
    pm.connectAttr(follicleShape.outTranslate, follicle.translate)

    follicleShape.parameterU.set(u)
    follicleShape.parameterV.set(v)

    return follicle

def findClosestUVCoordinate(geo, obj):
    closestPointOnMesh = pm.createNode("closestPointOnMesh")
    geoShape = geo.getShape()
    pm.connectAttr(geoShape.worldMesh, closestPointOnMesh.inMesh)
    pm.connectAttr(geoShape.worldMatrix, closestPointOnMesh.inputMatrix)
    loc = pm.spaceLocator(n=name.removeSuffix(geo.name())+'_LOC')
    pm.matchTransform(loc, obj)
    pm.connectAttr(loc.translate, closestPointOnMesh.inPosition)

    u = closestPointOnMesh.result.parameterU.get()
    v = closestPointOnMesh.result.parameterV.get()

    pm.delete(closestPointOnMesh, loc)

    return u, v

def makeControlFollowSkin(geo, ctrlTop):
    uv = findClosestUVCoordinate(geo, ctrlTop)
    prefix = name.removeSuffix(ctrlTop.name())
    follicle = createFollicle(geo, uv[0], uv[1], prefix)

    followGrp = pm.group(ctrlTop, n=prefix + 'Follow_GRP')
    compensateGrp = pm.group(ctrlTop, n=prefix + 'Compensate_GRP')

    pm.pointConstraint(follicle, followGrp, mo=True)

    multDivideNode = pm.createNode('multiplyDivide', n=prefix+'CompensateNode')
    pm.connectAttr(ctrlTop.translate, multDivideNode.input1)
    pm.connectAttr(multDivideNode.output, compensateGrp.translate)
    multDivideNode.input2X.set(-1)
    multDivideNode.input2Y.set(-1)
    multDivideNode.input2Z.set(-1)

    return follicle
<commit_msg>fix connection<commit_after>__author__ = 'Lorenzo Argentieri'

import pymel.core as pm
from maya import mel

from mayaLib.rigLib.utils import name

def createFollicle(geo, u, v, prefix):
    follicleShape = pm.createNode('follicle', n=prefix+'_FLCShape')
    geoShape = geo.getShape()
    follicle = follicleShape.getParent()
    follicle.rename(prefix + '_FLC')

    pm.connectAttr(geoShape.outMesh, follicleShape.inputMesh)
    pm.connectAttr(geoShape.worldMatrix, follicleShape.inputWorldMatrix)

    pm.connectAttr(follicleShape.outRotate, follicle.rotate)
    pm.connectAttr(follicleShape.outTranslate, follicle.translate)

    follicleShape.parameterU.set(u)
    follicleShape.parameterV.set(v)

    return follicle

def findClosestUVCoordinate(geo, obj):
    closestPointOnMesh = pm.createNode("closestPointOnMesh")
    geoShape = geo.getShape()
    pm.connectAttr(geoShape.worldMesh, closestPointOnMesh.inMesh)
    pm.connectAttr(geoShape.worldMatrix, closestPointOnMesh.inputMatrix)
    loc = pm.spaceLocator(n=name.removeSuffix(geo.name())+'_LOC')
    pm.matchTransform(loc, obj)
    pm.connectAttr(loc.translate, closestPointOnMesh.inPosition)

    u = closestPointOnMesh.result.parameterU.get()
    v = closestPointOnMesh.result.parameterV.get()

    pm.delete(closestPointOnMesh, loc)

    return u, v

def makeControlFollowSkin(geo, ctrlTop):
    uv = findClosestUVCoordinate(geo, ctrlTop)
    prefix = name.removeSuffix(ctrlTop.name())
    follicle = createFollicle(geo, uv[0], uv[1], prefix)

    followGrp = pm.group(ctrlTop, n=prefix + 'Follow_GRP')
    compensateGrp = pm.group(ctrlTop, n=prefix + 'Compensate_GRP')

    pm.pointConstraint(follicle, followGrp, mo=True)

    multDivideNode = pm.createNode('multiplyDivide', n=prefix+'CompensateNode')
    pm.connectAttr(followGrp.translate, multDivideNode.input1)
    pm.connectAttr(multDivideNode.output, compensateGrp.translate)
    multDivideNode.input2X.set(-1)
    multDivideNode.input2Y.set(-1)
    multDivideNode.input2Z.set(-1)

    return follicle
<|endoftext|>