<commit_before># Copyright (c) 2008 Duncan Fordyce
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
#  all copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import logging
import socket
import time

from oyoyo.parse import parse_raw_irc_command

class IRCClientError(Exception):
    pass

# Adapted from http://code.activestate.com/recipes/511490-implementation-of-the-token-bucket-algorithm/
class TokenBucket(object):
    """An implementation of the token bucket algorithm.

    >>> bucket = TokenBucket(80, 0.5)
    >>> bucket.consume(1)
    """
    def __init__(self, tokens, fill_rate):
        """tokens is the total tokens in the bucket. fill_rate is the
        rate in tokens/second that the bucket will be refilled."""
        self.capacity = float(tokens)
        self._tokens = float(tokens)
        self.fill_rate = float(fill_rate)
        self.timestamp = time.time()

    def consume(self, tokens):
        """Consume tokens from the bucket. Returns True if there were
        sufficient tokens otherwise False."""
        if tokens <= self.tokens:
            self._tokens -= tokens
        else:
            return False
        return True

    @property
    def tokens(self):
        now = time()
        if self._tokens < self.capacity:
            delta = self.fill_rate * (now - self.timestamp)
            self._tokens = min(self.capacity, self._tokens + delta)
        self.timestamp = now
        return self._tokens

def add_commands(d):
    def dec(cls):
        for c in d:
            def func(x):
                def gen(self, *a):
                    self.send(x.upper(), *a)
                return gen
            setattr(cls, c, func(c))
        return cls
    return dec
@add_commands(("join",
               "mode",
               "nick",
               "part",
               "kick",
               "who"))
class IRCClient(object):
    """ IRC Client class. This handles one connection to a server.
    This can be used either with or without IRCApp ( see connect() docs )
    """

    def __init__(self, cmd_handler, **kwargs):
        """ the first argument should be an object with attributes/methods named
        as the irc commands. You may subclass from one of the classes in
        oyoyo.cmdhandler for convenience but it is not required. The
        methods should have arguments (prefix, args). prefix is
        normally the sender of the command. args is a list of arguments.
        Its recommened you subclass oyoyo.cmdhandler.DefaultCommandHandler,
        this class provides defaults for callbacks that are required for
        normal IRC operation.

        all other arguments should be keyword arguments. The most commonly
        used will be nick, host and port. You can also specify an "on connect"
        callback. ( check the source for others )

        Warning: By default this class will not block on socket operations, this
        means if you use a plain while loop your app will consume 100% cpu.
        To enable blocking pass blocking=True.
        """

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.nickname = ""
        self.real_name = ""
        self.host = None
        self.port = None
        self.connect_cb = None
        self.blocking = True
        self.tokenbucket = TokenBucket(3, 1.63)

        self.__dict__.update(kwargs)
        self.command_handler = cmd_handler

        self._end = 0

    def send(self, *args, **kwargs):
        """ send a message to the connected server. all arguments are joined
        with a space for convenience, for example the following are identical

        >>> cli.send("JOIN " + some_room)
        >>> cli.send("JOIN", some_room)

        In python 2, all args must be of type str or unicode, *BUT* if they are
          unicode they will be converted to str with the encoding specified by
          the 'encoding' keyword argument (default 'utf8').
        In python 3, all args must be of type str or bytes, *BUT* if they are
          str they will be converted to bytes with the encoding specified by the
          'encoding' keyword argument (default 'utf8').
        """
        # Convert all args to bytes if not already
        encoding = kwargs.get('encoding') or 'utf_8'
        bargs = []
        for i,arg in enumerate(args):
            if isinstance(arg, str):
                bargs.append(bytes(arg, encoding))
            elif isinstance(arg, bytes):
                bargs.append(arg)
            elif arg is None:
                continue
            else:
                raise Exception(('Refusing to send arg at index {1} of the args from '+
                                 'provided: {0}').format(repr([(type(arg), arg)
                                                               for arg in args]), i))

        msg = bytes(" ", "utf_8").join(bargs)
        logging.info('---> send "{0}"'.format(msg))

        self.socket.send(msg + bytes("\r\n", "utf_8"))

    def connect(self):
        """ initiates the connection to the server set in self.host:self.port
        and returns a generator object.

        >>> cli = IRCClient(my_handler, host="irc.freenode.net", port=6667)
        >>> g = cli.connect()
        >>> while 1:
        ...     next(g)

        """
        try:
            logging.info('connecting to {0}:{1}'.format(self.host, self.port))
            retries = 0
            while True:
                try:
                    self.socket.connect(("{0}".format(self.host), self.port))
                    break
                except socket.error as e:
                    retries += 1
                    logging.warning('Error: {0}'.format(e))
                    if retries > 3:
                        break
            if not self.blocking:
                self.socket.setblocking(0)

            self.nick(self.nickname)
            self.user(self.nickname, self.real_name)

            if self.connect_cb:
                self.connect_cb(self)

            buffer = bytes()
            while not self._end:
                try:
                    buffer += self.socket.recv(1024)
                except socket.error as e:
                    if not self.blocking and e.errno == 11:
                        pass
                    else:
                        raise e
                else:
                    data = buffer.split(bytes("\n", "utf_8"))
                    buffer = data.pop()

                    for el in data:
                        prefix, command, args = parse_raw_irc_command(el)
                        logging.debug("processCommand ({2}){0}({1})".format(command,
                                                       [arg.decode('utf_8')
                                                        for arg in args
                                                        if isinstance(arg, bytes)], prefix))
                        try:
                            largs = list(args)
                            if prefix is not None:
                                prefix = prefix.decode("utf-8")
                            for i,arg in enumerate(largs):
                                if arg is not None: largs[i] = arg.decode('utf_8')
                            if command in self.command_handler:
                                self.command_handler[command](self, prefix,*largs)
                            elif "" in self.command_handler:
                                self.command_handler[""](self, prefix, command, *largs)
                        finally:
                            # error will of already been logged by the handler
                            pass

                yield True
        finally:
            if self.socket:
                logging.info('closing socket')
                self.socket.close()
    def msg(self, user, msg):
        for line in msg.split('\n'):
            while not self.tokenbucket.consume(1):
                time.sleep(1)
            self.send("PRIVMSG", user, ":{0}".format(line))
    privmsg = msg  # Same thing
    def notice(self, user, msg):
        for line in msg.split('\n'):
            self.send("NOTICE", user, ":{0}".format(line))
    def quit(self, msg):
        self.send("QUIT :{0}".format(msg))
    def identify(self, passwd, authuser="NickServ"):
        self.msg(authuser, "IDENTIFY {0}".format(passwd))
    def user(self, uname, rname):
        self.send("USER", uname, self.host, self.host,
                 rname or uname)
    def mainLoop(self):
        conn = self.connect()
        while True:
            next(conn)

<commit_msg>small bug<commit_after># Copyright (c) 2008 Duncan Fordyce
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
#  all copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import logging
import socket
import time

from oyoyo.parse import parse_raw_irc_command

class IRCClientError(Exception):
    pass

# Adapted from http://code.activestate.com/recipes/511490-implementation-of-the-token-bucket-algorithm/
class TokenBucket(object):
    """An implementation of the token bucket algorithm.

    >>> bucket = TokenBucket(80, 0.5)
    >>> bucket.consume(1)
    """
    def __init__(self, tokens, fill_rate):
        """tokens is the total tokens in the bucket. fill_rate is the
        rate in tokens/second that the bucket will be refilled."""
        self.capacity = float(tokens)
        self._tokens = float(tokens)
        self.fill_rate = float(fill_rate)
        self.timestamp = time.time()

    def consume(self, tokens):
        """Consume tokens from the bucket. Returns True if there were
        sufficient tokens otherwise False."""
        if tokens <= self.tokens:
            self._tokens -= tokens
        else:
            return False
        return True

    @property
    def tokens(self):
        now = time.time()
        if self._tokens < self.capacity:
            delta = self.fill_rate * (now - self.timestamp)
            self._tokens = min(self.capacity, self._tokens + delta)
        self.timestamp = now
        return self._tokens

def add_commands(d):
    def dec(cls):
        for c in d:
            def func(x):
                def gen(self, *a):
                    self.send(x.upper(), *a)
                return gen
            setattr(cls, c, func(c))
        return cls
    return dec
@add_commands(("join",
               "mode",
               "nick",
               "part",
               "kick",
               "who"))
class IRCClient(object):
    """ IRC Client class. This handles one connection to a server.
    This can be used either with or without IRCApp ( see connect() docs )
    """

    def __init__(self, cmd_handler, **kwargs):
        """ the first argument should be an object with attributes/methods named
        as the irc commands. You may subclass from one of the classes in
        oyoyo.cmdhandler for convenience but it is not required. The
        methods should have arguments (prefix, args). prefix is
        normally the sender of the command. args is a list of arguments.
        Its recommened you subclass oyoyo.cmdhandler.DefaultCommandHandler,
        this class provides defaults for callbacks that are required for
        normal IRC operation.

        all other arguments should be keyword arguments. The most commonly
        used will be nick, host and port. You can also specify an "on connect"
        callback. ( check the source for others )

        Warning: By default this class will not block on socket operations, this
        means if you use a plain while loop your app will consume 100% cpu.
        To enable blocking pass blocking=True.
        """

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.nickname = ""
        self.real_name = ""
        self.host = None
        self.port = None
        self.connect_cb = None
        self.blocking = True
        self.tokenbucket = TokenBucket(3, 1.63)

        self.__dict__.update(kwargs)
        self.command_handler = cmd_handler

        self._end = 0

    def send(self, *args, **kwargs):
        """ send a message to the connected server. all arguments are joined
        with a space for convenience, for example the following are identical

        >>> cli.send("JOIN " + some_room)
        >>> cli.send("JOIN", some_room)

        In python 2, all args must be of type str or unicode, *BUT* if they are
          unicode they will be converted to str with the encoding specified by
          the 'encoding' keyword argument (default 'utf8').
        In python 3, all args must be of type str or bytes, *BUT* if they are
          str they will be converted to bytes with the encoding specified by the
          'encoding' keyword argument (default 'utf8').
        """
        # Convert all args to bytes if not already
        encoding = kwargs.get('encoding') or 'utf_8'
        bargs = []
        for i,arg in enumerate(args):
            if isinstance(arg, str):
                bargs.append(bytes(arg, encoding))
            elif isinstance(arg, bytes):
                bargs.append(arg)
            elif arg is None:
                continue
            else:
                raise Exception(('Refusing to send arg at index {1} of the args from '+
                                 'provided: {0}').format(repr([(type(arg), arg)
                                                               for arg in args]), i))

        msg = bytes(" ", "utf_8").join(bargs)
        logging.info('---> send "{0}"'.format(msg))

        self.socket.send(msg + bytes("\r\n", "utf_8"))

    def connect(self):
        """ initiates the connection to the server set in self.host:self.port
        and returns a generator object.

        >>> cli = IRCClient(my_handler, host="irc.freenode.net", port=6667)
        >>> g = cli.connect()
        >>> while 1:
        ...     next(g)

        """
        try:
            logging.info('connecting to {0}:{1}'.format(self.host, self.port))
            retries = 0
            while True:
                try:
                    self.socket.connect(("{0}".format(self.host), self.port))
                    break
                except socket.error as e:
                    retries += 1
                    logging.warning('Error: {0}'.format(e))
                    if retries > 3:
                        break
            if not self.blocking:
                self.socket.setblocking(0)

            self.nick(self.nickname)
            self.user(self.nickname, self.real_name)

            if self.connect_cb:
                self.connect_cb(self)

            buffer = bytes()
            while not self._end:
                try:
                    buffer += self.socket.recv(1024)
                except socket.error as e:
                    if not self.blocking and e.errno == 11:
                        pass
                    else:
                        raise e
                else:
                    data = buffer.split(bytes("\n", "utf_8"))
                    buffer = data.pop()

                    for el in data:
                        prefix, command, args = parse_raw_irc_command(el)
                        logging.debug("processCommand ({2}){0}({1})".format(command,
                                                       [arg.decode('utf_8')
                                                        for arg in args
                                                        if isinstance(arg, bytes)], prefix))
                        try:
                            largs = list(args)
                            if prefix is not None:
                                prefix = prefix.decode("utf-8")
                            for i,arg in enumerate(largs):
                                if arg is not None: largs[i] = arg.decode('utf_8')
                            if command in self.command_handler:
                                self.command_handler[command](self, prefix,*largs)
                            elif "" in self.command_handler:
                                self.command_handler[""](self, prefix, command, *largs)
                        finally:
                            # error will of already been logged by the handler
                            pass

                yield True
        finally:
            if self.socket:
                logging.info('closing socket')
                self.socket.close()
    def msg(self, user, msg):
        for line in msg.split('\n'):
            while not self.tokenbucket.consume(1):
                time.sleep(1)
            self.send("PRIVMSG", user, ":{0}".format(line))
    privmsg = msg  # Same thing
    def notice(self, user, msg):
        for line in msg.split('\n'):
            self.send("NOTICE", user, ":{0}".format(line))
    def quit(self, msg):
        self.send("QUIT :{0}".format(msg))
    def identify(self, passwd, authuser="NickServ"):
        self.msg(authuser, "IDENTIFY {0}".format(passwd))
    def user(self, uname, rname):
        self.send("USER", uname, self.host, self.host,
                 rname or uname)
    def mainLoop(self):
        conn = self.connect()
        while True:
            next(conn)

<|endoftext|>


<commit_before># -*- coding: UTF-8 -*-
import logging
import models

from assopy.models import order_created, purchase_completed, ticket_for_user, user_created, user_identity_created
from conference.listeners import fare_price, fare_tickets
from conference.signals import attendees_connected
from conference.models import AttendeeProfile, Ticket
from django.conf import settings
from django.contrib.auth.models import User
from django.core.urlresolvers import reverse
from django.db.models.signals import post_save
from email_template import utils

log = logging.getLogger('p3')

def on_order_created(sender, **kwargs):
    if sender.total() == 0:
        on_purchase_completed(sender)
    elif sender.method == 'bank':
        utils.email(
            'bank-order-complete',
            ctx={'order': sender,},
            to=[sender.user.user.email]
        ).send()

    ritems = kwargs['raw_items']
    for fare, params in ritems:
        # se l'ordine contiene delle prenotazione alberghiere devo creare i
        # relativi ticket adesso, perchè le informazioni sul periodo le ho solo
        # tramite ritems
        if fare.code[0] == 'H':
            log.info(
                'The newly created order "%s" includes %d hotel reservations "%s" for the period: "%s" -> "%s".',
                sender.code,
                params['qty'],
                fare.code,
                params['period'][0],
                params['period'][1])
            loop = params['qty']
            if fare.code[1] == 'R':
                loop *= int(fare.code[2])
            for _ in range(loop):
                t = Ticket.objects.filter(fare=fare, user=sender.user.user, p3_conference_room=None)[0]
                room = models.TicketRoom(ticket=t)
                room.ticket_type = fare.code[1]
                room.room_type = models.HotelRoom.objects.get(conference=fare.conference, room_type='t%s' % fare.code[2])
                room.checkin = params['period'][0]
                room.checkout = params['period'][1]
                room.save()

order_created.connect(on_order_created)

def on_purchase_completed(sender, **kwargs):
    if sender.method == 'admin':
        return
    utils.email(
        'purchase-complete',
        ctx={
            'order': sender,
            'student': sender.orderitem_set.filter(ticket__fare__recipient_type='s').count() > 0,
        },
        to=[sender.user.user.email],
    ).send()

purchase_completed.connect(on_purchase_completed)

def on_ticket_for_user(sender, **kwargs):
    from p3 import dataaccess
    tickets = dataaccess.user_tickets(sender.user, settings.CONFERENCE_CONFERENCE, only_complete=True)
    map(kwargs['tickets'].append, tickets)

ticket_for_user.connect(on_ticket_for_user)

def on_profile_created(sender, **kw):
    if kw['created']:
        models.P3Profile(profile=kw['instance']).save()

post_save.connect(on_profile_created, sender=AttendeeProfile)

def on_user_created(sender, **kw):
    if kw['profile_complete']:
        # L'utente è stato creato utilizzando la form email+password, non avrò
        # altri dati
        AttendeeProfile.objects.getOrCreateForUser(sender.user)
    else:
        # L'utente ha utilizzato janrain, non creo il profilo ora ma aspetto la
        # prima identità
        pass

user_created.connect(on_user_created)

def on_user_identity_created(sender, **kw):
    identity = kw['identity']
    profile = AttendeeProfile.objects.getOrCreateForUser(sender.user)
    if identity.user.identities.count() > 1:
        # non è la prima identità non copio niente per non sovrascrivere
        # i dati cambiati dall'utente
        return

user_identity_created.connect(on_user_identity_created)

def calculate_hotel_reservation_price(sender, **kw):
    if sender.code[0] != 'H':
        return
    # il costo di una prenotazione alberghiera varia in funzione del periodo
    calc = kw['calc']
    period = calc['params']['period']
    room = models.HotelRoom.objects.get(conference=sender.conference, room_type='t' + sender.code[2])
    price = room.price(days=(period[1] - period[0]).days)
    if sender.code[1] == 'R':
        price *= int(sender.code[2])
    calc['total'] = price * calc['params']['qty']
fare_price.connect(calculate_hotel_reservation_price)

def create_hotel_tickets(sender, **kw):
    # solo per le prenotazioni di intere camere devo creare più biglietti, per
    # tutti gli altri casi mi va bene il comportamento standard
    if sender.code[:2] == 'HR':
        room_size = int(sender.code[2])
        for ix in range(room_size):
            t = Ticket(user=kw['params']['user'], fare=sender)
            t.fare_description = sender.name + (' (Occupant %s/%s)' % (ix+1, room_size))
            t.save()
            kw['params']['tickets'].append(t)

fare_tickets.connect(create_hotel_tickets)

# ridefinisco la user_tickets di assopy per includere i biglietti assegnati
from assopy import dataaccess as cd
_original = cd.user_tickets
def _user_tickets(u):
    data = _original(u)
    # aggiungo ai biglietti già individuati l'email della persona a cui sono
    # stati assegnati
    from p3.models import TicketConference
    tids = [ x['id'] for x in data ]
    info = dict([
        (x['ticket'], x['assigned_to'])
        for x in TicketConference.objects\
            .filter(ticket__in=tids)\
            .exclude(assigned_to="")\
            .values('ticket', 'assigned_to')])
    for x in data:
        try:
            x['note'] = 'to: %s' % info[x['id']]
        except KeyError:
            continue

    # aggiungo i biglietti assegnati all'utente
    qs = Ticket.objects\
        .filter(p3_conference__assigned_to__iexact=u.email)\
        .order_by('-fare__conference')\
        .select_related('fare')

    from assopy.models import Order
    def order(t):
        try:
            o = Order.objects\
                .get(orderitem__ticket=t)
        except models.Order.DoesNotExist:
            return {}
        return {
            'code': o.code,
            'created': o.created,
            'url': reverse('admin:assopy_order_change', args=(o.id,)),
        }
    for t in qs:
        data.append({
            'id': t.id,
            'type': t.ticket_type,
            'fare': {
                'code': t.fare.code,
                'conference': t.fare.conference,
                'recipient': t.fare.recipient_type,
            },
            'note': 'from: %s %s' % (t.user.first_name, t.user.last_name),
            'order': order(t),
        })
    return data
cd.user_tickets = _user_tickets

def _on_attendees_connected(sender, **kw):
    scanner = User.objects.get(id=kw['attendee1'])
    scanned = User.objects.get(id=kw['attendee2'])
    log.info(
        'User link: "%s %s" (%s) -> "%s %s" (%s)',
        scanner.first_name, scanner.last_name, scanner.id,
        scanned.first_name, scanned.last_name, scanned.id,
    )
    utils.email(
        'user-connected',
        ctx={
            'scanner': scanner,
            'scanned': scanned
        },
        to=[scanned.email]
    ).send()
attendees_connected.connect(_on_attendees_connected)
<commit_msg>sottoscrivo automaticamente i partecipanti del training al talk<commit_after># -*- coding: UTF-8 -*-
import logging
import models

from assopy.models import order_created, purchase_completed, ticket_for_user, user_created, user_identity_created
from conference.listeners import fare_price, fare_tickets
from conference.signals import attendees_connected, event_booked
from conference.models import AttendeeProfile, Ticket, Talk
from django.conf import settings
from django.contrib.auth.models import User
from django.core.urlresolvers import reverse
from django.db.models.signals import post_save
from email_template import utils

log = logging.getLogger('p3')

def on_order_created(sender, **kwargs):
    if sender.total() == 0:
        on_purchase_completed(sender)
    elif sender.method == 'bank':
        utils.email(
            'bank-order-complete',
            ctx={'order': sender,},
            to=[sender.user.user.email]
        ).send()

    ritems = kwargs['raw_items']
    for fare, params in ritems:
        # se l'ordine contiene delle prenotazione alberghiere devo creare i
        # relativi ticket adesso, perchè le informazioni sul periodo le ho solo
        # tramite ritems
        if fare.code[0] == 'H':
            log.info(
                'The newly created order "%s" includes %d hotel reservations "%s" for the period: "%s" -> "%s".',
                sender.code,
                params['qty'],
                fare.code,
                params['period'][0],
                params['period'][1])
            loop = params['qty']
            if fare.code[1] == 'R':
                loop *= int(fare.code[2])
            for _ in range(loop):
                t = Ticket.objects.filter(fare=fare, user=sender.user.user, p3_conference_room=None)[0]
                room = models.TicketRoom(ticket=t)
                room.ticket_type = fare.code[1]
                room.room_type = models.HotelRoom.objects.get(conference=fare.conference, room_type='t%s' % fare.code[2])
                room.checkin = params['period'][0]
                room.checkout = params['period'][1]
                room.save()

order_created.connect(on_order_created)

def on_purchase_completed(sender, **kwargs):
    if sender.method == 'admin':
        return
    utils.email(
        'purchase-complete',
        ctx={
            'order': sender,
            'student': sender.orderitem_set.filter(ticket__fare__recipient_type='s').count() > 0,
        },
        to=[sender.user.user.email],
    ).send()

purchase_completed.connect(on_purchase_completed)

def on_ticket_for_user(sender, **kwargs):
    from p3 import dataaccess
    tickets = dataaccess.user_tickets(sender.user, settings.CONFERENCE_CONFERENCE, only_complete=True)
    map(kwargs['tickets'].append, tickets)

ticket_for_user.connect(on_ticket_for_user)

def on_profile_created(sender, **kw):
    if kw['created']:
        models.P3Profile(profile=kw['instance']).save()

post_save.connect(on_profile_created, sender=AttendeeProfile)

def on_user_created(sender, **kw):
    if kw['profile_complete']:
        # L'utente è stato creato utilizzando la form email+password, non avrò
        # altri dati
        AttendeeProfile.objects.getOrCreateForUser(sender.user)
    else:
        # L'utente ha utilizzato janrain, non creo il profilo ora ma aspetto la
        # prima identità
        pass

user_created.connect(on_user_created)

def on_user_identity_created(sender, **kw):
    identity = kw['identity']
    profile = AttendeeProfile.objects.getOrCreateForUser(sender.user)
    if identity.user.identities.count() > 1:
        # non è la prima identità non copio niente per non sovrascrivere
        # i dati cambiati dall'utente
        return

user_identity_created.connect(on_user_identity_created)

def calculate_hotel_reservation_price(sender, **kw):
    if sender.code[0] != 'H':
        return
    # il costo di una prenotazione alberghiera varia in funzione del periodo
    calc = kw['calc']
    period = calc['params']['period']
    room = models.HotelRoom.objects.get(conference=sender.conference, room_type='t' + sender.code[2])
    price = room.price(days=(period[1] - period[0]).days)
    if sender.code[1] == 'R':
        price *= int(sender.code[2])
    calc['total'] = price * calc['params']['qty']
fare_price.connect(calculate_hotel_reservation_price)

def create_hotel_tickets(sender, **kw):
    # solo per le prenotazioni di intere camere devo creare più biglietti, per
    # tutti gli altri casi mi va bene il comportamento standard
    if sender.code[:2] == 'HR':
        room_size = int(sender.code[2])
        for ix in range(room_size):
            t = Ticket(user=kw['params']['user'], fare=sender)
            t.fare_description = sender.name + (' (Occupant %s/%s)' % (ix+1, room_size))
            t.save()
            kw['params']['tickets'].append(t)

fare_tickets.connect(create_hotel_tickets)

# ridefinisco la user_tickets di assopy per includere i biglietti assegnati
from assopy import dataaccess as cd
_original = cd.user_tickets
def _user_tickets(u):
    data = _original(u)
    # aggiungo ai biglietti già individuati l'email della persona a cui sono
    # stati assegnati
    from p3.models import TicketConference
    tids = [ x['id'] for x in data ]
    info = dict([
        (x['ticket'], x['assigned_to'])
        for x in TicketConference.objects\
            .filter(ticket__in=tids)\
            .exclude(assigned_to="")\
            .values('ticket', 'assigned_to')])
    for x in data:
        try:
            x['note'] = 'to: %s' % info[x['id']]
        except KeyError:
            continue

    # aggiungo i biglietti assegnati all'utente
    qs = Ticket.objects\
        .filter(p3_conference__assigned_to__iexact=u.email)\
        .order_by('-fare__conference')\
        .select_related('fare')

    from assopy.models import Order
    def order(t):
        try:
            o = Order.objects\
                .get(orderitem__ticket=t)
        except models.Order.DoesNotExist:
            return {}
        return {
            'code': o.code,
            'created': o.created,
            'url': reverse('admin:assopy_order_change', args=(o.id,)),
        }
    for t in qs:
        data.append({
            'id': t.id,
            'type': t.ticket_type,
            'fare': {
                'code': t.fare.code,
                'conference': t.fare.conference,
                'recipient': t.fare.recipient_type,
            },
            'note': 'from: %s %s' % (t.user.first_name, t.user.last_name),
            'order': order(t),
        })
    return data
cd.user_tickets = _user_tickets

def _on_attendees_connected(sender, **kw):
    scanner = User.objects.get(id=kw['attendee1'])
    scanned = User.objects.get(id=kw['attendee2'])
    log.info(
        'User link: "%s %s" (%s) -> "%s %s" (%s)',
        scanner.first_name, scanner.last_name, scanner.id,
        scanned.first_name, scanned.last_name, scanned.id,
    )
    utils.email(
        'user-connected',
        ctx={
            'scanner': scanner,
            'scanned': scanned
        },
        to=[scanned.email]
    ).send()
attendees_connected.connect(_on_attendees_connected)

def _on_event_booked(sender, **kw):
    from conference.dataaccess import event_data
    from hcomments.models import ThreadSubscription

    try:
        talk_id = event_data(kw['event_id'])['talk']['id']
    except Exception:
        return

    talk = Talk.objects.get(id=talk_id)
    user = User.objects.get(id=kw['user_id'])

    booked = kw['booked']
    if booked:
        log.info(
            "\"%s\" has booked the event \"%s\", automatically subscribed to the talk's comments",
            u'{0} {1}'.format(user.first_name, user.last_name), talk.title)
    else:
        log.info(
            "\"%s\" has cancelled the reservation for the event \"%s\", automatically unsubscribed to the talk's comments",
            u'{0} {1}'.format(user.first_name, user.last_name), talk.title)

    if booked:
        ThreadSubscription.objects.subscribe(talk, user)
    else:
        ThreadSubscription.objects.unsubscribe(talk, user)

event_booked.connect(_on_event_booked)
<|endoftext|>


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

# Copyright (c) 2013, Rethink Robotics
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. Neither the name of the Rethink Robotics nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""
Tool to tuck/untuck Baxter's arms to/from the shipping pose
"""
import math
import argparse
import threading

import roslib
roslib.load_manifest('tools')
import rospy
from std_msgs.msg import (
    Empty,
)
import baxter_interface
import dataflow

class Tuck(object):
    def __init__(self, tuck_cmd):
        self._limbs = ('left', 'right')
        self._arms = {
            'left': baxter_interface.Limb('left'),
            'right': baxter_interface.Limb('right'),
            }
        self._tuck = tuck_cmd
        self._tuck_rate = rospy.Rate(20.0) # Hz
        self._tuck_threshold = 0.6 # Radians
        self._tuck_state = {'left': False, 'right': False}
        self._joint_moves = {
            'tuck': {
                     'left': [ -1.0, -2.07,   3.0, 2.55,  0.0, 0.01,   0.0],
                     'right': [ 1.0, -2.07, -3.0, 2.55, -0.0, 0.01,  0.0]
                     },
            'untuck': {
                       'left': [-0.08,  -1.0, -1.19, 1.94, 0.67, 1.03, -0.50],
                       'right': [0.08,  -1.0, 1.19, 1.94, -0.67, 1.03, 0.50]
                       }
        }
        self._disable_pub = {
            'left': rospy.Publisher(
                 '/robot/limb/left/CollisionAvoidance/suppress_body_avoidance',
                 Empty),
            'right': rospy.Publisher(
                 '/robot/limb/right/CollisionAvoidance/suppress_body_avoidance',
                 Empty)
        }
        self._rs = baxter_interface.RobotEnable()
        self._supervised_tuck()

    def _check_tucked(self, goal):
        for limb in self._limbs:
            diff = 0.0
            for joint_idx, joint in enumerate(self._arms[limb].joint_names()):
                if joint[-2:] in ('s0', 's1', 'e0') or goal == 'untuck':
                    diff = diff + abs(self._joint_moves[goal][limb][joint_idx] - self._arms[limb].joint_angle(joint))
            self._tuck_state[limb] = diff < self._tuck_threshold

    def _reset_tucked_state(self):
        for limb in self._limbs:
            self._tuck_state[limb] = False

    def _move_to(self, tuck, disabled):
        self._rs.enable()
        while any(self._tuck_state[limb] != True for limb in self._limbs) and not rospy.is_shutdown():
            for limb in self._limbs:
                if disabled:
                    self._disable_pub[limb].publish(Empty())
                self._arms[limb].set_joint_positions(dict(zip(self._arms[limb].joint_names(),self._joint_moves[tuck][limb])))
            self._check_tucked(tuck)
            self._tuck_rate.sleep()
        self._reset_tucked_state()
        if (tuck == 'tuck'):
            self._rs.disable()
        return

    def _supervised_tuck(self):
        # Check our starting state to see if arms are tucked
        self._check_tucked('tuck')
        # Tuck Arms
        if self._tuck == True:
            # If arms are already tucked, report this to user and exit.
            if all(self._tuck_state[limb] == True for limb in self._limbs):
                rospy.loginfo("Arms already in 'Tucked' position.")
                return
            else:
                # Move to neutral location before tucking arms to avoid damaging the robot
                self._move_to('untuck', False)
                # Disable collision and Tuck Arms
                self._move_to('tuck', True)
                return
        # Untuck Arms
        else:
            # If arms are tucked disable collision and untuck arms.
            if all(self._tuck_state[limb] == True for limb in self._limbs):
                self._reset_tucked_state()
                # Disable collision and untuck Arms
                self._move_to('untuck', True)
                return
            # If arms already untucked, move to neutral location
            else:
                self._move_to('untuck', False)
                return

def main(tuck):
    print("Initializing node... ")
    rospy.init_node("tuck_arms")
    print("Moving head to neutral position")
    head = baxter_interface.Head()
    head.set_pan(0.0, 5.0)
    print("%sucking arms" % ("T" if tuck else "Unt",))
    Tuck(tuck)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    tuck_group = parser.add_mutually_exclusive_group(required=True)
    tuck_group.add_argument("-t", "--tuck", dest="tuck", \
        action='store_true', default=False, help="tuck arms")
    tuck_group.add_argument("-u", "--untuck", dest="untuck", \
        action='store_true', default=False, help="untuck arms")
    args = parser.parse_args()
    main(args.tuck)
<commit_msg>Updates to handle limbs in the opposite tucked/untucked states.<commit_after>#!/usr/bin/env python

# Copyright (c) 2013, Rethink Robotics
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. Neither the name of the Rethink Robotics nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""
Tool to tuck/untuck Baxter's arms to/from the shipping pose
"""
import math
import argparse
import threading

import roslib
roslib.load_manifest('tools')
import rospy
from std_msgs.msg import (
    Empty,
)
import baxter_interface
import dataflow

class Tuck(object):
    def __init__(self, tuck_cmd):
        self._limbs = ('left', 'right')
        self._arms = {
            'left': baxter_interface.Limb('left'),
            'right': baxter_interface.Limb('right'),
            }
        self._tuck = tuck_cmd
        self._tuck_rate = rospy.Rate(20.0) # Hz
        self._tuck_threshold = 0.6 # Radians
        self._tuck_state = {'left': False, 'right': False}
        self._joint_moves = {
            'tuck': {
                     'left': [ -1.0, -2.07,   3.0, 2.55,  0.0, 0.01,   0.0],
                     'right': [ 1.0, -2.07, -3.0, 2.55, -0.0, 0.01,  0.0]
                     },
            'untuck': {
                       'left': [-0.08,  -1.0, -1.19, 1.94, 0.67, 1.03, -0.50],
                       'right': [0.08,  -1.0, 1.19, 1.94, -0.67, 1.03, 0.50]
                       }
        }
        self._disable_pub = {
            'left': rospy.Publisher(
                 '/robot/limb/left/CollisionAvoidance/suppress_body_avoidance',
                 Empty),
            'right': rospy.Publisher(
                 '/robot/limb/right/CollisionAvoidance/suppress_body_avoidance',
                 Empty)
        }
        self._rs = baxter_interface.RobotEnable()
        self._supervised_tuck()

    def _check_tucked(self, goal):
        for limb in self._limbs:
            diff = 0.0
            for joint_idx, joint in enumerate(self._arms[limb].joint_names()):
                if joint[-2:] in ('s0', 's1', 'e0') or goal == 'untuck':
                    diff = diff + abs(self._joint_moves[goal][limb][joint_idx] - self._arms[limb].joint_angle(joint))
            self._tuck_state[limb] = diff < self._tuck_threshold

    def _reset_tucked_state(self):
        for limb in self._limbs:
            self._tuck_state[limb] = False

    def _move_to(self, tuck, disabled):
        self._rs.enable()
        while any(self._tuck_state[limb] != True for limb in self._limbs) and not rospy.is_shutdown():
            for limb in self._limbs:
                if disabled:
                    self._disable_pub[limb].publish(Empty())
                self._arms[limb].set_joint_positions(dict(zip(self._arms[limb].joint_names(),self._joint_moves[tuck][limb])))
            self._check_tucked(tuck)
            self._tuck_rate.sleep()
        self._reset_tucked_state()
        if (tuck == 'tuck'):
            self._rs.disable()
        return

    def _supervised_tuck(self):
        # Check our starting state to see if arms are tucked
        self._check_tucked('tuck')
        # Tuck Arms
        if self._tuck == True:
            # If arms are already tucked, report this to user and exit.
            if all(self._tuck_state[limb] == True for limb in self._limbs):
                rospy.loginfo("Arms already in 'Tucked' position.")
                return
            else:
                any_tucked = any(self._tuck_state[limb] == True for limb in self._limbs)
                # Move to neutral location before tucking arms to avoid damaging the robot
                self._move_to('untuck', any_tucked)
                # Disable collision and Tuck Arms
                self._move_to('tuck', True)
                return
        # Untuck Arms
        else:
            # If arms are tucked disable collision and untuck arms.
            if any(self._tuck_state[limb] == True for limb in self._limbs):
                self._reset_tucked_state()
                # Disable collision and untuck Arms
                self._move_to('untuck', True)
                return
            # If arms already untucked, move to neutral location
            else:
                self._move_to('untuck', False)
                return

def main(tuck):
    print("Initializing node... ")
    rospy.init_node("tuck_arms")
    print("Moving head to neutral position")
    head = baxter_interface.Head()
    head.set_pan(0.0, 5.0)
    print("%sucking arms" % ("T" if tuck else "Unt",))
    Tuck(tuck)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    tuck_group = parser.add_mutually_exclusive_group(required=True)
    tuck_group.add_argument("-t", "--tuck", dest="tuck", \
        action='store_true', default=False, help="tuck arms")
    tuck_group.add_argument("-u", "--untuck", dest="untuck", \
        action='store_true', default=False, help="untuck arms")
    args = parser.parse_args()
    main(args.tuck)
<|endoftext|>


<commit_before>#!/usr/bin/python

import sys, os, subprocess, shutil

# Parse configuration

make_exe = False
make_pkg = False
make_msi = False
target = None

for arg in sys.argv:
    if arg == "--exe":
        make_exe = True
    elif arg == "--pkg":
        make_pkg = True
    elif arg == "--msi":
        make_msi = True
    elif "--target" in arg:
        target = arg.split("=")[1]

print
print "target: " + str(target)
print "exe: " + str(make_exe)
print "pkg: " + str(make_pkg)
print "msi: " + str(make_msi)
print

if target == None:
    print "specify --target"
    sys.exit(1)

def run(args):
    retval = subprocess.call(args)
    if retval != 0:
        print "call failed: " + str(args)
        sys.exit(1)

INPUT_DIR = "./in"
OUTPUT_DIR = "./out"
TEMP_DIR = "./tmp"
RUSTC_PACKAGE_NAME = "rust"
COMBINED_PACKAGE_NAME = "rust-combined"

# Create the temp directory
if os.path.isdir(TEMP_DIR):
    shutil.rmtree(TEMP_DIR)
os.mkdir(TEMP_DIR)

if not os.path.isdir(OUTPUT_DIR):
    os.mkdir(OUTPUT_DIR)

# The names of the packages that need to be combined via rust-installer
components = [RUSTC_PACKAGE_NAME, "cargo", "rust-docs"]
if "pc-windows-gnu" in target:
    components += ["rust-mingw"]

# Now find the names of the tarballs that belong to those components
inputs = []
version = None
rustc_installer = None
cargo_installer = None
docs_installer = None
mingw_installer = None
source_tarball = None
for component in components:
    component_installer = None
    for filename in os.listdir(INPUT_DIR):
        if target in filename and component in filename:
            # Hack: several components contain 'rust' in the name
            if not (component == "rust" and ("rust-docs" in filename or "rust-mingw" in filename)):
                component_installer = filename

        if "-src" in filename:
            source_tarball = filename

    if not component_installer:
        print "unable to find installer for component " + component + ", target " + target
        sys.exit(1)
    inputs += [INPUT_DIR + "/" + component_installer]

    # Extract the version from the filename
    if component == RUSTC_PACKAGE_NAME:
        s = component_installer[len(RUSTC_PACKAGE_NAME) + 1:]
        p = s.find(target)
        version = s[:p - 1]
        rustc_installer = component_installer
    if component == "cargo":
        cargo_installer = component_installer
    if component == "rust-docs":
        docs_installer = component_installer
    if component == "rust-mingw":
        mingw_installer = component_installer

assert version != None
assert rustc_installer != None
assert source_tarball != None

# Set up the overlay of license info
run(["tar", "xzf", INPUT_DIR + "/" + rustc_installer, "-C", TEMP_DIR, ])

rustc_dir = TEMP_DIR + "/" + rustc_installer[:len(rustc_installer) - len(".tar.gz")]

overlay_dir = TEMP_DIR + "/overlay"
os.mkdir(overlay_dir)
shutil.copyfile(rustc_dir + "/COPYRIGHT", overlay_dir + "/COPYRIGHT")
shutil.copyfile(rustc_dir + "/LICENSE-APACHE", overlay_dir + "/LICENSE-APACHE")
shutil.copyfile(rustc_dir + "/LICENSE-MIT", overlay_dir + "/LICENSE-MIT")
shutil.copyfile(rustc_dir + "/README.md", overlay_dir + "/README.md")

# Combine the installers
tarball_list=",".join(inputs)
package_name=COMBINED_PACKAGE_NAME + "-" + version + "-" + target
run(["sh", "./rust-installer/combine-installers.sh",
     "--product-name=Rust",
     "--verify-bin=rustc",
     "--rel-manifest-dir=rustlib",
     "--success-message=Rust-is-ready-to-roll",
     "--work-dir=" + TEMP_DIR + "/work",
     "--output-dir=" + OUTPUT_DIR,
     "--package-name=" + package_name,
     "--legacy-manifest-dirs=rustlib,cargo",
     "--input-tarballs=" + tarball_list,
     "--non-installed-overlay=" + overlay_dir
 ])

# Everything below here is used for producing non-rust-installer packaging

# Create the LICENSE.txt file used in some GUI installers
license_file = TEMP_DIR + "/LICENSE.txt"
cmd = "cat " + \
    rustc_dir + "/COPYRIGHT " + rustc_dir + "/LICENSE-APACHE " + rustc_dir + "/LICENSE-MIT " + \
    "> " + license_file
run(["sh", "-c", cmd])

# Fish out the following Makefile variables from the source code or rebuild them somehow.
# Currently these are needed by the Windows installer.
CFG_RELEASE_NUM = None
CFG_BETA_CYCLE = None
CFG_RELEASE = None
CFG_PACKAGE_NAME = None
CFG_BUILD = None
CFG_PACKAGE_VERS = None

run(["tar", "xzf", INPUT_DIR + "/" + source_tarball, "-C", TEMP_DIR])
source_dir = os.path.join(TEMP_DIR, source_tarball.replace("-src.tar.gz", ""))

for line in open(source_dir + "/mk/main.mk"):
    if "CFG_RELEASE_NUM" in line and CFG_RELEASE_NUM == None:
        CFG_RELEASE_NUM = line.split("=")[1].strip()
        assert len(CFG_RELEASE_NUM) > 0
    if "CFG_BETA_CYCLE" in line and CFG_BETA_CYCLE == None:
        CFG_BETA_CYCLE = line.split("=")[1].strip()
        # NB: This can be an empty string

assert CFG_RELEASE_NUM != None

# FIXME Temporary hack
if CFG_BETA_CYCLE == None:
    CFG_BETA_CYCLE = ""
assert CFG_BETA_CYCLE != None

# Guess the channel from the source tarball
channel = None
if "nightly" in source_tarball:
    channel = "nightly"
elif "beta" in source_tarball:
    channel = "beta"
elif "dev" in source_tarball:
    channel = "dev"
else:
    channel = "stable"

# Logic reproduced from main.mk
if channel == "stable":
    CFG_RELEASE=CFG_RELEASE_NUM
    CFG_PACKAGE_VERS=CFG_RELEASE_NUM
elif channel == "beta":
    CFG_RELEASE=CFG_RELEASE_NUM + "-beta" + CFG_BETA_CYCLE
    CFG_PACKAGE_VERS="beta"
elif channel == "nightly":
    CFG_RELEASE=CFG_RELEASE_NUM + "-nightly"
    CFG_PACKAGE_VERS="nightly"
elif channel == "dev":
    CFG_RELEASE=CFG_RELEASE_NUM + "-dev"
    CFG_PACKAGE_VERS=CFG_RELEASE_NUM + "-dev"
else:
    print "unknown release channel"
    sys.exit(1)

# This should be the same as the name on the tarballs
CFG_PACKAGE_NAME=COMBINED_PACKAGE_NAME + CFG_PACKAGE_VERS
CFG_BUILD=target

os.environ["CFG_RELEASE_NUM"] = CFG_RELEASE_NUM
os.environ["CFG_RELEASE"] = CFG_RELEASE
os.environ["CFG_PACKAGE_NAME"] = CFG_PACKAGE_NAME
os.environ["CFG_BUILD"] = CFG_BUILD

if make_pkg:
    print "creating .pkg"
    os.mkdir(TEMP_DIR + "/pkg")
    shutil.copytree(TEMP_DIR + "/work/" + package_name, TEMP_DIR + "/pkg/root")
    os.mkdir(TEMP_DIR + "/pkg/res")
    shutil.copyfile(TEMP_DIR + "/LICENSE.txt", TEMP_DIR + "/pkg/res/LICENSE.txt")
    shutil.copyfile("./pkg/welcome.rtf", TEMP_DIR + "/pkg/res/welcome.rtf")
    shutil.copyfile("./gfx/rust-logo.png", TEMP_DIR + "/pkg/res/rust-logo.png")
    pkgbuild_cmd = "pkgbuild --identifier org.rust-lang.rust " + \
        "--root " + TEMP_DIR + "/pkg/root " + TEMP_DIR + "/pkg/rust.pkg"
    run(["sh", "-c", pkgbuild_cmd])
    productbuild_cmd = "productbuild --distribution ./pkg/Distribution.xml " + \
        "--resources " + TEMP_DIR + "/pkg/res " + OUTPUT_DIR + "/" + package_name + ".pkg " + \
        "--package-path " + TEMP_DIR + "/pkg"
    run(["sh", "-c", productbuild_cmd])

if make_exe:
    print "creating .exe"
    assert docs_installer != None
    assert mingw_installer != None
    assert cargo_installer != None
    exe_temp_dir = TEMP_DIR + "/exe"
    os.mkdir(exe_temp_dir)
    run(["tar", "xzf", INPUT_DIR + "/" + rustc_installer, "-C", exe_temp_dir])
    run(["tar", "xzf", INPUT_DIR + "/" + docs_installer, "-C", exe_temp_dir])
    run(["tar", "xzf", INPUT_DIR + "/" + mingw_installer, "-C", exe_temp_dir])
    run(["tar", "xzf", INPUT_DIR + "/" + cargo_installer, "-C", exe_temp_dir])
    orig_rustc_dir = exe_temp_dir + "/" + rustc_installer.replace(".tar.gz", "")
    orig_docs_dir = exe_temp_dir + "/" + docs_installer.replace(".tar.gz", "")
    orig_mingw_dir = exe_temp_dir + "/" + mingw_installer.replace(".tar.gz", "")
    orig_cargo_dir = exe_temp_dir + "/" + cargo_installer.replace(".tar.gz", "")

    # Move these to locations needed by the iscc script
    rustc_dir = exe_temp_dir + "/rustc"
    docs_dir = exe_temp_dir + "/rust-docs"
    mingw_dir = exe_temp_dir + "/rust-mingw"
    cargo_dir = exe_temp_dir + "/cargo"
    os.rename(orig_rustc_dir, rustc_dir)
    os.rename(orig_docs_dir, docs_dir)
    os.rename(orig_mingw_dir, mingw_dir)
    os.rename(orig_cargo_dir, cargo_dir)

    # Remove the installer files we don't need
    dir_comp_pairs = [(rustc_dir, "rustc"), (docs_dir, "rust-docs"),
                      (mingw_dir, "rust-mingw"), (cargo_dir, "cargo")]
    for dir_and_component in dir_comp_pairs:
        dir_ = dir_and_component[0]
        component = dir_and_component[1]
        for file_ in ["components", "install.sh", "rust-installer-version"]:
            os.remove(dir_ + "/" + file_)
        os.remove(dir_ + "/manifest-" + component + ".in")

    # Copy installer files, etc.
    shutil.copyfile("./exe/rust.iss", exe_temp_dir + "/rust.iss")
    shutil.copyfile("./exe/modpath.iss", exe_temp_dir + "/modpath.iss")
    shutil.copyfile("./exe/upgrade.iss", exe_temp_dir + "/upgrade.iss")
    shutil.copyfile("./gfx/rust-logo.ico", exe_temp_dir + "/rust-logo.ico")
    shutil.copyfile(TEMP_DIR + "/LICENSE.txt", exe_temp_dir + "/LICENSE.txt")

    cwd=os.getcwd()
    os.chdir(exe_temp_dir)
    run(["iscc", "rust.iss"])
    os.chdir(cwd)

    exefile = CFG_PACKAGE_NAME + "-" + CFG_BUILD + ".exe"
    shutil.move(exe_temp_dir + "/" + exefile, OUTPUT_DIR + "/" + exefile)

# TODO Produce .msi
<commit_msg>Fix exe package name<commit_after>#!/usr/bin/python

import sys, os, subprocess, shutil

# Parse configuration

make_exe = False
make_pkg = False
make_msi = False
target = None

for arg in sys.argv:
    if arg == "--exe":
        make_exe = True
    elif arg == "--pkg":
        make_pkg = True
    elif arg == "--msi":
        make_msi = True
    elif "--target" in arg:
        target = arg.split("=")[1]

print
print "target: " + str(target)
print "exe: " + str(make_exe)
print "pkg: " + str(make_pkg)
print "msi: " + str(make_msi)
print

if target == None:
    print "specify --target"
    sys.exit(1)

def run(args):
    retval = subprocess.call(args)
    if retval != 0:
        print "call failed: " + str(args)
        sys.exit(1)

INPUT_DIR = "./in"
OUTPUT_DIR = "./out"
TEMP_DIR = "./tmp"
RUSTC_PACKAGE_NAME = "rust"
COMBINED_PACKAGE_NAME = "rust-combined"

# Create the temp directory
if os.path.isdir(TEMP_DIR):
    shutil.rmtree(TEMP_DIR)
os.mkdir(TEMP_DIR)

if not os.path.isdir(OUTPUT_DIR):
    os.mkdir(OUTPUT_DIR)

# The names of the packages that need to be combined via rust-installer
components = [RUSTC_PACKAGE_NAME, "cargo", "rust-docs"]
if "pc-windows-gnu" in target:
    components += ["rust-mingw"]

# Now find the names of the tarballs that belong to those components
inputs = []
version = None
rustc_installer = None
cargo_installer = None
docs_installer = None
mingw_installer = None
source_tarball = None
for component in components:
    component_installer = None
    for filename in os.listdir(INPUT_DIR):
        if target in filename and component in filename:
            # Hack: several components contain 'rust' in the name
            if not (component == "rust" and ("rust-docs" in filename or "rust-mingw" in filename)):
                component_installer = filename

        if "-src" in filename:
            source_tarball = filename

    if not component_installer:
        print "unable to find installer for component " + component + ", target " + target
        sys.exit(1)
    inputs += [INPUT_DIR + "/" + component_installer]

    # Extract the version from the filename
    if component == RUSTC_PACKAGE_NAME:
        s = component_installer[len(RUSTC_PACKAGE_NAME) + 1:]
        p = s.find(target)
        version = s[:p - 1]
        rustc_installer = component_installer
    if component == "cargo":
        cargo_installer = component_installer
    if component == "rust-docs":
        docs_installer = component_installer
    if component == "rust-mingw":
        mingw_installer = component_installer

assert version != None
assert rustc_installer != None
assert source_tarball != None

# Set up the overlay of license info
run(["tar", "xzf", INPUT_DIR + "/" + rustc_installer, "-C", TEMP_DIR, ])

rustc_dir = TEMP_DIR + "/" + rustc_installer[:len(rustc_installer) - len(".tar.gz")]

overlay_dir = TEMP_DIR + "/overlay"
os.mkdir(overlay_dir)
shutil.copyfile(rustc_dir + "/COPYRIGHT", overlay_dir + "/COPYRIGHT")
shutil.copyfile(rustc_dir + "/LICENSE-APACHE", overlay_dir + "/LICENSE-APACHE")
shutil.copyfile(rustc_dir + "/LICENSE-MIT", overlay_dir + "/LICENSE-MIT")
shutil.copyfile(rustc_dir + "/README.md", overlay_dir + "/README.md")

# Combine the installers
tarball_list=",".join(inputs)
package_name=COMBINED_PACKAGE_NAME + "-" + version + "-" + target
run(["sh", "./rust-installer/combine-installers.sh",
     "--product-name=Rust",
     "--verify-bin=rustc",
     "--rel-manifest-dir=rustlib",
     "--success-message=Rust-is-ready-to-roll",
     "--work-dir=" + TEMP_DIR + "/work",
     "--output-dir=" + OUTPUT_DIR,
     "--package-name=" + package_name,
     "--legacy-manifest-dirs=rustlib,cargo",
     "--input-tarballs=" + tarball_list,
     "--non-installed-overlay=" + overlay_dir
 ])

# Everything below here is used for producing non-rust-installer packaging

# Create the LICENSE.txt file used in some GUI installers
license_file = TEMP_DIR + "/LICENSE.txt"
cmd = "cat " + \
    rustc_dir + "/COPYRIGHT " + rustc_dir + "/LICENSE-APACHE " + rustc_dir + "/LICENSE-MIT " + \
    "> " + license_file
run(["sh", "-c", cmd])

# Fish out the following Makefile variables from the source code or rebuild them somehow.
# Currently these are needed by the Windows installer.
CFG_RELEASE_NUM = None
CFG_BETA_CYCLE = None
CFG_RELEASE = None
CFG_PACKAGE_NAME = None
CFG_BUILD = None
CFG_PACKAGE_VERS = None

run(["tar", "xzf", INPUT_DIR + "/" + source_tarball, "-C", TEMP_DIR])
source_dir = os.path.join(TEMP_DIR, source_tarball.replace("-src.tar.gz", ""))

for line in open(source_dir + "/mk/main.mk"):
    if "CFG_RELEASE_NUM" in line and CFG_RELEASE_NUM == None:
        CFG_RELEASE_NUM = line.split("=")[1].strip()
        assert len(CFG_RELEASE_NUM) > 0
    if "CFG_BETA_CYCLE" in line and CFG_BETA_CYCLE == None:
        CFG_BETA_CYCLE = line.split("=")[1].strip()
        # NB: This can be an empty string

assert CFG_RELEASE_NUM != None

# FIXME Temporary hack
if CFG_BETA_CYCLE == None:
    CFG_BETA_CYCLE = ""
assert CFG_BETA_CYCLE != None

# Guess the channel from the source tarball
channel = None
if "nightly" in source_tarball:
    channel = "nightly"
elif "beta" in source_tarball:
    channel = "beta"
elif "dev" in source_tarball:
    channel = "dev"
else:
    channel = "stable"

# Logic reproduced from main.mk
if channel == "stable":
    CFG_RELEASE=CFG_RELEASE_NUM
    CFG_PACKAGE_VERS=CFG_RELEASE_NUM
elif channel == "beta":
    CFG_RELEASE=CFG_RELEASE_NUM + "-beta" + CFG_BETA_CYCLE
    CFG_PACKAGE_VERS="beta"
elif channel == "nightly":
    CFG_RELEASE=CFG_RELEASE_NUM + "-nightly"
    CFG_PACKAGE_VERS="nightly"
elif channel == "dev":
    CFG_RELEASE=CFG_RELEASE_NUM + "-dev"
    CFG_PACKAGE_VERS=CFG_RELEASE_NUM + "-dev"
else:
    print "unknown release channel"
    sys.exit(1)

# This should be the same as the name on the tarballs
CFG_PACKAGE_NAME=COMBINED_PACKAGE_NAME + "-" + CFG_PACKAGE_VERS
CFG_BUILD=target

os.environ["CFG_RELEASE_NUM"] = CFG_RELEASE_NUM
os.environ["CFG_RELEASE"] = CFG_RELEASE
os.environ["CFG_PACKAGE_NAME"] = CFG_PACKAGE_NAME
os.environ["CFG_BUILD"] = CFG_BUILD

if make_pkg:
    print "creating .pkg"
    os.mkdir(TEMP_DIR + "/pkg")
    shutil.copytree(TEMP_DIR + "/work/" + package_name, TEMP_DIR + "/pkg/root")
    os.mkdir(TEMP_DIR + "/pkg/res")
    shutil.copyfile(TEMP_DIR + "/LICENSE.txt", TEMP_DIR + "/pkg/res/LICENSE.txt")
    shutil.copyfile("./pkg/welcome.rtf", TEMP_DIR + "/pkg/res/welcome.rtf")
    shutil.copyfile("./gfx/rust-logo.png", TEMP_DIR + "/pkg/res/rust-logo.png")
    pkgbuild_cmd = "pkgbuild --identifier org.rust-lang.rust " + \
        "--root " + TEMP_DIR + "/pkg/root " + TEMP_DIR + "/pkg/rust.pkg"
    run(["sh", "-c", pkgbuild_cmd])
    productbuild_cmd = "productbuild --distribution ./pkg/Distribution.xml " + \
        "--resources " + TEMP_DIR + "/pkg/res " + OUTPUT_DIR + "/" + package_name + ".pkg " + \
        "--package-path " + TEMP_DIR + "/pkg"
    run(["sh", "-c", productbuild_cmd])

if make_exe:
    print "creating .exe"
    assert docs_installer != None
    assert mingw_installer != None
    assert cargo_installer != None
    exe_temp_dir = TEMP_DIR + "/exe"
    os.mkdir(exe_temp_dir)
    run(["tar", "xzf", INPUT_DIR + "/" + rustc_installer, "-C", exe_temp_dir])
    run(["tar", "xzf", INPUT_DIR + "/" + docs_installer, "-C", exe_temp_dir])
    run(["tar", "xzf", INPUT_DIR + "/" + mingw_installer, "-C", exe_temp_dir])
    run(["tar", "xzf", INPUT_DIR + "/" + cargo_installer, "-C", exe_temp_dir])
    orig_rustc_dir = exe_temp_dir + "/" + rustc_installer.replace(".tar.gz", "")
    orig_docs_dir = exe_temp_dir + "/" + docs_installer.replace(".tar.gz", "")
    orig_mingw_dir = exe_temp_dir + "/" + mingw_installer.replace(".tar.gz", "")
    orig_cargo_dir = exe_temp_dir + "/" + cargo_installer.replace(".tar.gz", "")

    # Move these to locations needed by the iscc script
    rustc_dir = exe_temp_dir + "/rustc"
    docs_dir = exe_temp_dir + "/rust-docs"
    mingw_dir = exe_temp_dir + "/rust-mingw"
    cargo_dir = exe_temp_dir + "/cargo"
    os.rename(orig_rustc_dir, rustc_dir)
    os.rename(orig_docs_dir, docs_dir)
    os.rename(orig_mingw_dir, mingw_dir)
    os.rename(orig_cargo_dir, cargo_dir)

    # Remove the installer files we don't need
    dir_comp_pairs = [(rustc_dir, "rustc"), (docs_dir, "rust-docs"),
                      (mingw_dir, "rust-mingw"), (cargo_dir, "cargo")]
    for dir_and_component in dir_comp_pairs:
        dir_ = dir_and_component[0]
        component = dir_and_component[1]
        for file_ in ["components", "install.sh", "rust-installer-version"]:
            os.remove(dir_ + "/" + file_)
        os.remove(dir_ + "/manifest-" + component + ".in")

    # Copy installer files, etc.
    shutil.copyfile("./exe/rust.iss", exe_temp_dir + "/rust.iss")
    shutil.copyfile("./exe/modpath.iss", exe_temp_dir + "/modpath.iss")
    shutil.copyfile("./exe/upgrade.iss", exe_temp_dir + "/upgrade.iss")
    shutil.copyfile("./gfx/rust-logo.ico", exe_temp_dir + "/rust-logo.ico")
    shutil.copyfile(TEMP_DIR + "/LICENSE.txt", exe_temp_dir + "/LICENSE.txt")

    cwd=os.getcwd()
    os.chdir(exe_temp_dir)
    run(["iscc", "rust.iss"])
    os.chdir(cwd)

    exefile = CFG_PACKAGE_NAME + "-" + CFG_BUILD + ".exe"
    shutil.move(exe_temp_dir + "/" + exefile, OUTPUT_DIR + "/" + exefile)

# TODO Produce .msi
<|endoftext|>


<commit_before>import io
import os
import xml.etree.ElementTree as ET

from ..fileio.datareaders import load_data, _language_column_names

class BaseModel(object):
    """
    Base class from which all substitution model classes are descended.
    Implements generic functionality which is common to all substitution
    models, such as rate variation.
    """

    def __init__(self, model_config, global_config):
        """
        Parse configuration options, load data from file and pre-process data.
        """
        self.messages = []
        self.config = global_config

        self.name = model_config["name"]
        self.data_filename = model_config["data"]
        self.clock = model_config.get("clock", "")
        self.features = model_config.get("features",["*"])
        self.constant_feature = False
        self.frequencies = model_config.get("frequencies", "empirical")
        self.pruned = model_config.get("pruned", False)
        self.rate_variation = model_config.get("rate_variation", False)
        self.remove_constant_features = model_config.get("remove_constant_features", True)
        self.minimum_data = float(model_config.get("minimum_data", 0))
        self.lang_column = model_config.get("language_column", None)
        self.substitution_name = self.__class__.__name__
        self.data_separator = ","

        self.data = load_data(self.data_filename, file_format=model_config.get("file_format",None), lang_column=model_config.get("language_column",None))
        self.build_feature_filter()
        self.process()

    def build_feature_filter(self):
        """
        Create the self.feature_filter attribute, which is a set of feature
        names that functions analogously to Configuration.lang_filter
        attribute.
        """
        if self.features == ["*"]:
            random_iso = list(self.data.keys())[0]
            self.features = list(self.data[random_iso].keys())
            # Need to remove the language ID column
            if self.lang_column:
                self.features.remove(self.lang_column)
            else:
                # If no language column name was explicitly given, just
                # remove the first of the automatically-recognised names
                # which we encounter:
                for lc in _language_column_names:
                    if lc in self.features:
                        self.features.remove(lc)
                        break
        self.feature_filter = set(self.features)

    def process(self):
        """
        Subsample the data set to include only those languages and features
        which are compatible with the settings.
        """
        self.apply_language_filter()
        self.apply_feature_filter()
        self.compute_feature_properties()
        self.remove_unwanted_features()
        if self.pruned:
            self.messages.append("""[DEPENDENCY] Model %s: Pruned trees are implemented in the BEAST package "BEASTlabs".""" % self.name)

    def apply_language_filter(self):
        """
        Remove all languages from the data set which are not part of the
        configured language filter.
        """
        all_langs = self.data.keys()
        langs_to_remove = [l for l in all_langs if not self.config.filter_language(l)]
        for lang in langs_to_remove:
            self.data.pop(lang)
        # Make sure we've not removed all languages
        if not self.data.keys():
            raise ValueError("Language filters leave nothing in the dataset for model '%s'!" % self.name)
        # Keep a sorted list so that the order of things in XML is deterministic
        self.languages = sorted(list(self.data.keys()))

    def apply_feature_filter(self):
        """
        Remove all features from the data set which are not part of the
        configured feature filter.
        """
        self.features = set()
        for lang in self.languages:
            features_in_data = set(self.data[lang].keys())
            features_to_keep = features_in_data & self.feature_filter
            self.features |= features_to_keep
            features_to_remove = features_in_data - features_to_keep
            for feat in features_to_remove:
                self.data[lang].pop(feat)
        self.features = sorted(list(self.features))

    def compute_feature_properties(self):
        """
        Compute various items of metadata for all remaining features.
        """

        self.valuecounts = {}
        self.unique_values = {}
        self.missing_ratios = {}
        self.counts = {}
        self.dimensions = {}
        self.codemaps = {}
        for f in self.features:
            # Compute various things
            all_values = [self.data[l].get(f,"?") for l in self.data]
            missing_data_ratio = all_values.count("?") / (1.0*len(all_values))
            non_q_values = [v for v in all_values if v != "?"]
            counts = {}
            for v in non_q_values:
                counts[v] = non_q_values.count(v)
            unique_values = list(set(non_q_values))
            # Sort unique_values carefully.
            # Possibly all feature values are numeric strings, e.g. "1", "2", "3".
            # If we sort these as strings then we get weird things like "10" < "2".
            # This can actually matter for things like ordinal models.
            # So convert these to ints first...
            if all([v.isdigit() for v in unique_values]):
                unique_values = list(map(int, unique_values))
                unique_values.sort()
                unique_values = list(map(str, unique_values))
            # ...otherwise, just sort normally
            else:
                unique_values.sort()
            self.unique_values[f] = unique_values

            N = len(unique_values)
            self.valuecounts[f] = N
            self.missing_ratios[f] = missing_data_ratio
            self.counts[f] = counts
            self.dimensions[f] = N*(N-1) // 2
            self.codemaps[f] = self.build_codemap(unique_values)

    def remove_unwanted_features(self):
        """
        Remove any undesirable features from the dataset, such as those with
        no data for the configured set of languages, constant features, etc.
        """

        bad_feats = []
        for f in self.features:

            # Exclude features with no data
            if self.valuecounts[f] == 0:
                self.messages.append("""[INFO] Model "%s": Feature %s excluded because there are no datapoints for selected languages.""" % (self.name, f))
                bad_feats.append(f)
                continue

            # Exclude features with lots of missing data
            missing_ratio = self.missing_ratios[f]
            if int(100*(1.0-missing_ratio)) < self.minimum_data:
                self.messages.append("""[INFO] Model "%s": Feature %s excluded because of excessive missing data (%d%%).""" % (self.name, f, int(missing_ratio*100)))
                bad_feats.append(f)
                continue

            # Exclude constant features
            if self.valuecounts[f] == 1:
                if self.remove_constant_features:
                    self.messages.append("""[INFO] Model "%s": Feature %s excluded because its value is constant across selected languages.  Set "remove_constant_features=False" in config to stop this.""" % (self.name, f))
                    bad_feats.append(f)
                    continue
                else:
                    self.constant_feature = True

        for bad in bad_feats:
            self.features.remove(bad)
            for lang in self.languages:
                if bad in self.data[lang]:
                    self.data[lang].pop(bad)

        # Make sure there's something left
        if not self.features:
            raise ValueError("No features specified for model %s!" % self.name)
        self.features.sort()
        self.messages.append("""[INFO] Model "%s": Using %d features from data source %s""" % (self.name, len(self.features), self.data_filename))
        if self.constant_feature and self.rate_variation:
            self.messages.append("""[WARNING] Model "%s": Rate variation enabled with constant features retained in data.  This may skew rate estimates for non-constant features.""" % self.name)

    def build_codemap(self, unique_values):
        """
        Build a codemap string for a feature.
        """
        N = len(unique_values)
        codemapbits = []
        codemapbits.append(",".join(["%s=%d" % (v,n) for (n,v) in enumerate(unique_values)]))
        codemapbits.append("?=" + " ".join([str(n) for n in range(0,N)]))
        codemapbits.append("-=" + " ".join([str(n) for n in range(0,N)]))
        return ",".join(codemapbits)

    def add_misc(self, beast):
        pass

    def add_state(self, state):
        """
        Add parameters for Gamma-distributed rate heterogenetiy, if
        configured.
        """
        if self.rate_variation:
            plate = ET.SubElement(state, "plate", {
                "var":"feature",
                "range":",".join(self.features)})
            param = ET.SubElement(plate, "parameter", {
                "id":"featureClockRate:%s:$(feature)" % self.name,
                "name":"stateNode"})
            param.text="1.0"

            parameter = ET.SubElement(state, "parameter", {"id":"featureClockRateGammaShape:%s" % self.name, "lower":"0.0","upper":"100.0","name":"stateNode"})
            parameter.text="2.0"
            parameter = ET.SubElement(state, "parameter", {"id":"featureClockRateGammaScale:%s" % self.name, "name":"stateNode"})
            parameter.text="0.5"

    def add_prior(self, prior):
        """
        Add prior distributions for Gamma-distributed rate heterogenetiy, if
        configured.
        """
        if self.rate_variation:
            sub_prior = ET.SubElement(prior, "prior", {"id":"featureClockRatePrior.s:%s" % self.name, "name":"distribution"})
            compound = ET.SubElement(sub_prior, "input", {"id":"featureClockRateCompound:%s" % self.name, "spec":"beast.core.parameter.CompoundValuable", "name":"x"})
            plate = ET.SubElement(compound, "plate", {
                "var":"feature",
                "range":",".join(self.features)})
            ET.SubElement(plate, "var", {
                "idref":"featureClockRate:%s:$(feature)" % self.name})
            gamma  = ET.SubElement(sub_prior, "input", {"id":"featureClockRatePriorGamma:%s" % self.name, "spec":"beast.math.distributions.Gamma", "name":"distr", "alpha":"@featureClockRateGammaShape:%s" % self.name, "beta":"@featureClockRateGammaScale:%s" % self.name})

            sub_prior = ET.SubElement(prior, "prior", {"id":"featureClockRateGammaShapePrior.s:%s" % self.name, "name":"distribution", "x":"@featureClockRateGammaShape:%s" % self.name})
            ET.SubElement(sub_prior, "Exponential", {"id":"featureClockRateGammaShapePriorExponential.s:%s" % self.name, "mean":"1.0", "name":"distr"})

    def add_likelihood(self, likelihood):
        """
        Add likelihood distribution corresponding to all features in the
        dataset.
        """
        for n, f in enumerate(self.features):
            fname = "%s:%s" % (self.name, f)
            attribs = {"id":"featureLikelihood:%s" % fname,"spec":"TreeLikelihood","useAmbiguities":"true"}
            if self.pruned:
                distribution = ET.SubElement(likelihood, "distribution",attribs)
                # Create pruned tree
                tree_id = "Tree.t:prunedBeastlingTree.%s" % fname
                tree = ET.SubElement(distribution, "tree", {"id":tree_id, "spec":"beast.evolution.tree.PrunedTree","quickshortcut":"true","assert":"false"})
                ET.SubElement(tree, "tree", {"idref":"Tree.t:beastlingTree"})
                ET.SubElement(tree, "alignment", {"idref":"pruned_data_%s"%fname})
                # Create pruned branchrate
                self.clock.add_pruned_branchrate_model(distribution, fname, tree_id)
            else:
                attribs["branchRateModel"] = "@%s" % self.clock.branchrate_model_id
                attribs["tree"] = "@Tree.t:beastlingTree"
                distribution = ET.SubElement(likelihood, "distribution",attribs)

            # Sitemodel
            self.add_sitemodel(distribution, f, fname)

            # Data
            self.add_feature_data(distribution, n, f, fname)

    def add_sitemodel(self, beast):
        pass

    def add_master_data(self, beast):
        self.filters = {}
        data = ET.SubElement(beast, "data", {
            "id":"data_%s" % self.name,
            "name":"data_%s" % self.name,
            "dataType":"integer"})
        for lang in self.languages:
            formatted_points = [self.format_datapoint(f, self.data[lang][f]) for f in self.features]
            value_string = self.data_separator.join(formatted_points)
            self.filters = {f:str(n+1) for n, f in enumerate(self.features)}
            seq = ET.SubElement(data, "sequence", {
                "id":"data_%s:%s" % (self.name, lang),
                "taxon":lang,
                "value":value_string})

    def format_datapoint(self, feature, point):
        return str(point)

    def add_feature_data(self, distribution, index, feature, fname):
        """
        Add <data> element corresponding to the indicated feature, descending
        from the indicated likelihood distribution.
        """
        if self.pruned:
            pruned_align = ET.SubElement(distribution,"data",{"id":"pruned_data_%s" % fname, "spec":"PrunedAlignment"})
            parent = pruned_align
            name = "source"
        else:
            parent = distribution
            name = "data"
        data = ET.SubElement(parent, name, {
            "id":"data_%s" % fname,
            "spec":"FilteredAlignment",
            "data":"@data_%s" % self.name,
            "filter":self.filters[feature]})
        data.append(self.get_userdatatype(feature, fname))
        return data

    def get_userdatatype(self, feature, fname):
        return ET.Element("userDataType", {"id":"featureDataType.%s"%fname,"spec":"beast.evolution.datatype.UserDataType","codeMap":self.codemaps[feature],"codelength":"-1","states":str(self.valuecounts[feature])})

    def add_operators(self, run):
        """
        Add <operators> for individual feature substitution rates if rate
        variation is configured.
        """
        if self.rate_variation:
            # UpDownOperator to scale the Gamma distribution for this model's
            # feature rates
            updown = ET.SubElement(run, "operator", {"id":"featureClockRateGammaUpDown:%s" % self.name, "spec":"UpDownOperator", "scaleFactor":"0.5","weight":"0.3"})
            ET.SubElement(updown, "parameter", {"idref":"featureClockRateGammaShape:%s" % self.name, "name":"up"})
            ET.SubElement(updown, "parameter", {"idref":"featureClockRateGammaScale:%s" % self.name, "name":"down"})

    def add_param_logs(self, logger):
        """
        Add entires to the logfile corresponding to individual feature
        substition rates if rate variation is configured.
        """
        if self.config.log_fine_probs:
            plate = ET.SubElement(logger, "plate", {
                "var":"feature",
                "range":",".join(self.features)})
            ET.SubElement(plate, "log", {
                "idref":"featureLikelihood:%s:$(feature)" % self.name})
            if self.rate_variation:
                ET.SubElement(logger,"log",{"idref":"featureClockRatePrior.s:%s" % self.name})
                ET.SubElement(logger,"log",{"idref":"featureClockRateGammaShapePrior.s:%s" % self.name})

        if self.rate_variation:
            plate = ET.SubElement(logger, "plate", {
                "var":"feature",
                "range":",".join(self.features)})
            ET.SubElement(plate, "log", {
                "idref":"featureClockRate:%s:$(feature)" % self.name})
            # Log the shape, but not the scale, as it is always 1 / shape
            ET.SubElement(logger,"log",{"idref":"featureClockRateGammaShape:%s" % self.name})
<commit_msg>Yet another pass at correct calculation of FilteredAlignment filter values.<commit_after>import io
import os
import xml.etree.ElementTree as ET

from ..fileio.datareaders import load_data, _language_column_names

class BaseModel(object):
    """
    Base class from which all substitution model classes are descended.
    Implements generic functionality which is common to all substitution
    models, such as rate variation.
    """

    def __init__(self, model_config, global_config):
        """
        Parse configuration options, load data from file and pre-process data.
        """
        self.messages = []
        self.config = global_config

        self.name = model_config["name"]
        self.data_filename = model_config["data"]
        self.clock = model_config.get("clock", "")
        self.features = model_config.get("features",["*"])
        self.constant_feature = False
        self.frequencies = model_config.get("frequencies", "empirical")
        self.pruned = model_config.get("pruned", False)
        self.rate_variation = model_config.get("rate_variation", False)
        self.remove_constant_features = model_config.get("remove_constant_features", True)
        self.minimum_data = float(model_config.get("minimum_data", 0))
        self.lang_column = model_config.get("language_column", None)
        self.substitution_name = self.__class__.__name__
        self.data_separator = ","

        self.data = load_data(self.data_filename, file_format=model_config.get("file_format",None), lang_column=model_config.get("language_column",None))
        self.build_feature_filter()
        self.process()

    def build_feature_filter(self):
        """
        Create the self.feature_filter attribute, which is a set of feature
        names that functions analogously to Configuration.lang_filter
        attribute.
        """
        if self.features == ["*"]:
            random_iso = list(self.data.keys())[0]
            self.features = list(self.data[random_iso].keys())
            # Need to remove the language ID column
            if self.lang_column:
                self.features.remove(self.lang_column)
            else:
                # If no language column name was explicitly given, just
                # remove the first of the automatically-recognised names
                # which we encounter:
                for lc in _language_column_names:
                    if lc in self.features:
                        self.features.remove(lc)
                        break
        self.feature_filter = set(self.features)

    def process(self):
        """
        Subsample the data set to include only those languages and features
        which are compatible with the settings.
        """
        self.apply_language_filter()
        self.apply_feature_filter()
        self.compute_feature_properties()
        self.remove_unwanted_features()
        if self.pruned:
            self.messages.append("""[DEPENDENCY] Model %s: Pruned trees are implemented in the BEAST package "BEASTlabs".""" % self.name)

    def apply_language_filter(self):
        """
        Remove all languages from the data set which are not part of the
        configured language filter.
        """
        all_langs = self.data.keys()
        langs_to_remove = [l for l in all_langs if not self.config.filter_language(l)]
        for lang in langs_to_remove:
            self.data.pop(lang)
        # Make sure we've not removed all languages
        if not self.data.keys():
            raise ValueError("Language filters leave nothing in the dataset for model '%s'!" % self.name)
        # Keep a sorted list so that the order of things in XML is deterministic
        self.languages = sorted(list(self.data.keys()))

    def apply_feature_filter(self):
        """
        Remove all features from the data set which are not part of the
        configured feature filter.
        """
        self.features = set()
        for lang in self.languages:
            features_in_data = set(self.data[lang].keys())
            features_to_keep = features_in_data & self.feature_filter
            self.features |= features_to_keep
            features_to_remove = features_in_data - features_to_keep
            for feat in features_to_remove:
                self.data[lang].pop(feat)
        self.features = sorted(list(self.features))

    def compute_feature_properties(self):
        """
        Compute various items of metadata for all remaining features.
        """

        self.valuecounts = {}
        self.unique_values = {}
        self.missing_ratios = {}
        self.counts = {}
        self.dimensions = {}
        self.codemaps = {}
        for f in self.features:
            # Compute various things
            all_values = [self.data[l].get(f,"?") for l in self.data]
            missing_data_ratio = all_values.count("?") / (1.0*len(all_values))
            non_q_values = [v for v in all_values if v != "?"]
            counts = {}
            for v in non_q_values:
                counts[v] = non_q_values.count(v)
            unique_values = list(set(non_q_values))
            # Sort unique_values carefully.
            # Possibly all feature values are numeric strings, e.g. "1", "2", "3".
            # If we sort these as strings then we get weird things like "10" < "2".
            # This can actually matter for things like ordinal models.
            # So convert these to ints first...
            if all([v.isdigit() for v in unique_values]):
                unique_values = list(map(int, unique_values))
                unique_values.sort()
                unique_values = list(map(str, unique_values))
            # ...otherwise, just sort normally
            else:
                unique_values.sort()
            self.unique_values[f] = unique_values

            N = len(unique_values)
            self.valuecounts[f] = N
            self.missing_ratios[f] = missing_data_ratio
            self.counts[f] = counts
            self.dimensions[f] = N*(N-1) // 2
            self.codemaps[f] = self.build_codemap(unique_values)

    def remove_unwanted_features(self):
        """
        Remove any undesirable features from the dataset, such as those with
        no data for the configured set of languages, constant features, etc.
        """

        bad_feats = []
        for f in self.features:

            # Exclude features with no data
            if self.valuecounts[f] == 0:
                self.messages.append("""[INFO] Model "%s": Feature %s excluded because there are no datapoints for selected languages.""" % (self.name, f))
                bad_feats.append(f)
                continue

            # Exclude features with lots of missing data
            missing_ratio = self.missing_ratios[f]
            if int(100*(1.0-missing_ratio)) < self.minimum_data:
                self.messages.append("""[INFO] Model "%s": Feature %s excluded because of excessive missing data (%d%%).""" % (self.name, f, int(missing_ratio*100)))
                bad_feats.append(f)
                continue

            # Exclude constant features
            if self.valuecounts[f] == 1:
                if self.remove_constant_features:
                    self.messages.append("""[INFO] Model "%s": Feature %s excluded because its value is constant across selected languages.  Set "remove_constant_features=False" in config to stop this.""" % (self.name, f))
                    bad_feats.append(f)
                    continue
                else:
                    self.constant_feature = True

        for bad in bad_feats:
            self.features.remove(bad)
            for lang in self.languages:
                if bad in self.data[lang]:
                    self.data[lang].pop(bad)

        # Make sure there's something left
        if not self.features:
            raise ValueError("No features specified for model %s!" % self.name)
        self.features.sort()
        self.messages.append("""[INFO] Model "%s": Using %d features from data source %s""" % (self.name, len(self.features), self.data_filename))
        if self.constant_feature and self.rate_variation:
            self.messages.append("""[WARNING] Model "%s": Rate variation enabled with constant features retained in data.  This may skew rate estimates for non-constant features.""" % self.name)

    def build_codemap(self, unique_values):
        """
        Build a codemap string for a feature.
        """
        N = len(unique_values)
        codemapbits = []
        codemapbits.append(",".join(["%s=%d" % (v,n) for (n,v) in enumerate(unique_values)]))
        codemapbits.append("?=" + " ".join([str(n) for n in range(0,N)]))
        codemapbits.append("-=" + " ".join([str(n) for n in range(0,N)]))
        return ",".join(codemapbits)

    def add_misc(self, beast):
        pass

    def add_state(self, state):
        """
        Add parameters for Gamma-distributed rate heterogenetiy, if
        configured.
        """
        if self.rate_variation:
            plate = ET.SubElement(state, "plate", {
                "var":"feature",
                "range":",".join(self.features)})
            param = ET.SubElement(plate, "parameter", {
                "id":"featureClockRate:%s:$(feature)" % self.name,
                "name":"stateNode"})
            param.text="1.0"

            parameter = ET.SubElement(state, "parameter", {"id":"featureClockRateGammaShape:%s" % self.name, "lower":"0.0","upper":"100.0","name":"stateNode"})
            parameter.text="2.0"
            parameter = ET.SubElement(state, "parameter", {"id":"featureClockRateGammaScale:%s" % self.name, "name":"stateNode"})
            parameter.text="0.5"

    def add_prior(self, prior):
        """
        Add prior distributions for Gamma-distributed rate heterogenetiy, if
        configured.
        """
        if self.rate_variation:
            sub_prior = ET.SubElement(prior, "prior", {"id":"featureClockRatePrior.s:%s" % self.name, "name":"distribution"})
            compound = ET.SubElement(sub_prior, "input", {"id":"featureClockRateCompound:%s" % self.name, "spec":"beast.core.parameter.CompoundValuable", "name":"x"})
            plate = ET.SubElement(compound, "plate", {
                "var":"feature",
                "range":",".join(self.features)})
            ET.SubElement(plate, "var", {
                "idref":"featureClockRate:%s:$(feature)" % self.name})
            gamma  = ET.SubElement(sub_prior, "input", {"id":"featureClockRatePriorGamma:%s" % self.name, "spec":"beast.math.distributions.Gamma", "name":"distr", "alpha":"@featureClockRateGammaShape:%s" % self.name, "beta":"@featureClockRateGammaScale:%s" % self.name})

            sub_prior = ET.SubElement(prior, "prior", {"id":"featureClockRateGammaShapePrior.s:%s" % self.name, "name":"distribution", "x":"@featureClockRateGammaShape:%s" % self.name})
            ET.SubElement(sub_prior, "Exponential", {"id":"featureClockRateGammaShapePriorExponential.s:%s" % self.name, "mean":"1.0", "name":"distr"})

    def add_likelihood(self, likelihood):
        """
        Add likelihood distribution corresponding to all features in the
        dataset.
        """
        for n, f in enumerate(self.features):
            fname = "%s:%s" % (self.name, f)
            attribs = {"id":"featureLikelihood:%s" % fname,"spec":"TreeLikelihood","useAmbiguities":"true"}
            if self.pruned:
                distribution = ET.SubElement(likelihood, "distribution",attribs)
                # Create pruned tree
                tree_id = "Tree.t:prunedBeastlingTree.%s" % fname
                tree = ET.SubElement(distribution, "tree", {"id":tree_id, "spec":"beast.evolution.tree.PrunedTree","quickshortcut":"true","assert":"false"})
                ET.SubElement(tree, "tree", {"idref":"Tree.t:beastlingTree"})
                ET.SubElement(tree, "alignment", {"idref":"pruned_data_%s"%fname})
                # Create pruned branchrate
                self.clock.add_pruned_branchrate_model(distribution, fname, tree_id)
            else:
                attribs["branchRateModel"] = "@%s" % self.clock.branchrate_model_id
                attribs["tree"] = "@Tree.t:beastlingTree"
                distribution = ET.SubElement(likelihood, "distribution",attribs)

            # Sitemodel
            self.add_sitemodel(distribution, f, fname)

            # Data
            self.add_feature_data(distribution, n, f, fname)

    def add_sitemodel(self, beast):
        pass

    def add_master_data(self, beast):
        self.filters = {}
        data = ET.SubElement(beast, "data", {
            "id":"data_%s" % self.name,
            "name":"data_%s" % self.name,
            "dataType":"integer"})
        for lang in self.languages:
            formatted_points = [self.format_datapoint(f, self.data[lang][f]) for f in self.features]
            value_string = self.data_separator.join(formatted_points)
            if not self.filters:
                n = 1
                for f, x in zip(self.features, formatted_points):
                    # Find out how many sites in the sequence correspond to this feature
                    if self.data_separator:
                        length = len(x.split(self.data_separator))
                    else:
                        length = len(x)
                    # Format the FilteredAlignment filter appropriately
                    if length == 1:
                        self.filters[f] = str(n)
                    else:
                        self.filters[f] = "%d-%d" % (n, n+length-1)
                    n += length
            seq = ET.SubElement(data, "sequence", {
                "id":"data_%s:%s" % (self.name, lang),
                "taxon":lang,
                "value":value_string})

    def format_datapoint(self, feature, point):
        return str(point)

    def add_feature_data(self, distribution, index, feature, fname):
        """
        Add <data> element corresponding to the indicated feature, descending
        from the indicated likelihood distribution.
        """
        if self.pruned:
            pruned_align = ET.SubElement(distribution,"data",{"id":"pruned_data_%s" % fname, "spec":"PrunedAlignment"})
            parent = pruned_align
            name = "source"
        else:
            parent = distribution
            name = "data"
        data = ET.SubElement(parent, name, {
            "id":"data_%s" % fname,
            "spec":"FilteredAlignment",
            "data":"@data_%s" % self.name,
            "filter":self.filters[feature]})
        data.append(self.get_userdatatype(feature, fname))
        return data

    def get_userdatatype(self, feature, fname):
        return ET.Element("userDataType", {"id":"featureDataType.%s"%fname,"spec":"beast.evolution.datatype.UserDataType","codeMap":self.codemaps[feature],"codelength":"-1","states":str(self.valuecounts[feature])})

    def add_operators(self, run):
        """
        Add <operators> for individual feature substitution rates if rate
        variation is configured.
        """
        if self.rate_variation:
            # UpDownOperator to scale the Gamma distribution for this model's
            # feature rates
            updown = ET.SubElement(run, "operator", {"id":"featureClockRateGammaUpDown:%s" % self.name, "spec":"UpDownOperator", "scaleFactor":"0.5","weight":"0.3"})
            ET.SubElement(updown, "parameter", {"idref":"featureClockRateGammaShape:%s" % self.name, "name":"up"})
            ET.SubElement(updown, "parameter", {"idref":"featureClockRateGammaScale:%s" % self.name, "name":"down"})

    def add_param_logs(self, logger):
        """
        Add entires to the logfile corresponding to individual feature
        substition rates if rate variation is configured.
        """
        if self.config.log_fine_probs:
            plate = ET.SubElement(logger, "plate", {
                "var":"feature",
                "range":",".join(self.features)})
            ET.SubElement(plate, "log", {
                "idref":"featureLikelihood:%s:$(feature)" % self.name})
            if self.rate_variation:
                ET.SubElement(logger,"log",{"idref":"featureClockRatePrior.s:%s" % self.name})
                ET.SubElement(logger,"log",{"idref":"featureClockRateGammaShapePrior.s:%s" % self.name})

        if self.rate_variation:
            plate = ET.SubElement(logger, "plate", {
                "var":"feature",
                "range":",".join(self.features)})
            ET.SubElement(plate, "log", {
                "idref":"featureClockRate:%s:$(feature)" % self.name})
            # Log the shape, but not the scale, as it is always 1 / shape
            ET.SubElement(logger,"log",{"idref":"featureClockRateGammaShape:%s" % self.name})
<|endoftext|>