<commit_before>from decimal import Decimal

from dna.vendor import refprop as rp

class Error(Exception):
    """Base class for exceptions in this module."""
    pass

class InputError(Error):
    """Exception raised for errors in the input.

    Attributes:
        expr -- input expression in which the error occurred
        msg  -- explanation of the error
    """

    def __init__(self, expr, msg):
        self.expr = expr
        self.msg = msg

rp.setup('def', 'ammonia', 'water')
rp.setref(hrf='def',ixflag=1)
molWNH3 = rp.info(1)['wmm']
molWH2O = rp.info(2)['wmm']

def state(node):

    use_rp = True

    if 'media' in node:
        usermedia = {
            'hitecxl': {'cp': 1.447, 'tmin': 130, 'tmax': 490},
            'hitec': {'cp': 1.5617, 'tmin': 180, 'tmax': 560}
        }

        if node['media'] in usermedia:
            node['cp'] = usermedia[node['media']]['cp']
            node['tmin'] = usermedia[node['media']]['tmin']
            node['tmax'] = usermedia[node['media']]['tmax']
            use_rp = False

    if use_rp:
        return refpropState(node)
    else:
        return cpBasedState(node)

def molarToMass(x):
    mass_nh3 = float(x)*molWNH3;
    mass_h2o = (1-float(x))*molWH2O;

    mass_total = mass_nh3 + mass_h2o

    return [mass_nh3/mass_total, 1-mass_nh3/mass_total]

def massToMolar(y):
    m_nh3 = float(y)/molWNH3
    m_h2o = (1-float(y))/molWH2O

    mt = m_nh3 + m_h2o

    return [m_nh3/mt, 1-m_nh3/mt]

def toRefprop(node):
    prop = {}

    x = massToMolar(node['y'])

    molWmix = float(x[0]) * float(molWNH3) + float(x[1]) * float(molWH2O) # g/mol

    if 'mdot' in node:
        prop['mdot'] = node['mdot']

    if 'p' in node:
        prop['p'] = node['p']*100 # hPa > kPa

    if 'e' in node:
        prop['e'] = node['e']*molWmix # kJ/kg > J/mol

    if 'h' in node:
        prop['h'] = node['h']*molWmix # kJ/kg*K > J/mol*K

    if 'D' in node:
        prop['D'] = node['D']/molWmix*1000 # kg/L > mol/m3

    if 's' in node:
        prop['s'] = node['s']*molWmix # kJ/kg*K > J/mol*K

    if 't' in node:
        prop['t'] = node['t'] + 273.15 # C > K

    if 'cv' in node:
        prop['cv'] = node['cv']*molWmix # kJ/kg*K > J/mol*K

    if 'cp' in node:
        prop['cp'] = node['cp']*molWmix # kJ/kg*K > J/mol*K

    if 'q' in node:
        prop['q'] = massToMolar(node['q'])[0] # kg/kg > mol/mol

    if 'y' in node:
        prop['x'] = massToMolar(node['y']) # kg/kg > mol/mol

    if 'yvap' in node:
        prop['xvap'] = massToMolar(node['yvap']) # kg/kg > mol/mol

    if 'yliq' in node:
        prop['xliq'] = massToMolar(node['yliq']) # kg/kg > mol/mol

    return prop

def fromRefprop(prop):
    node = {}

    x = prop['x']

    molWmix = float(x[0]) * float(molWNH3) + float(x[1]) * float(molWH2O)

    if 'mdot' in prop:
        node['mdot'] = prop['mdot']

    # Convert mol to kg, K to C, kPa to hPa
    if 'p' in prop:
        node['p'] = prop['p']/100 # kPa > hPa

    if 'e' in prop:
        node['e'] = prop['e']/molWmix # J/mol > kJ/kg

    if 'h' in prop:
        node['h'] = prop['h']/molWmix # J/mol > kJ/kg

    if 'D' in prop:
        node['D'] = prop['D']*molWmix/1000 # mol/L > kg/m3

    if 's' in prop:
        node['s'] = prop['s']/molWmix # J/mol*K > kJ/kg*K

    if 't' in prop:
        node['t'] = prop['t'] - 273.15 # K > C

    if 'cv' in prop:
        node['cv'] = prop['cv']/molWmix # J/mol*K > kJ/kg*K

    if 'cp' in prop:
        node['cp'] = prop['cp']/molWmix # J/mol*K > kJ/kg*K

    if 'q' in prop:
        node['q'] = molarToMass(prop['q'])[0] # mol/mol > kg/kg

    if 'x' in prop:
        node['y'] = molarToMass(prop['x'][0])[0] # mol/mol > kg/kg

    if 'xvap' in prop:
        node['yvap'] = molarToMass(prop['xvap'][0])[0] # mol/mol > kg/kg

    if 'xliq' in prop:
        node['yliq'] = molarToMass(prop['xliq'][0])[0] # mol/mol > kg/kg

    return node

def _PHworkaround(_node, depth = 1):

    if depth > 5:
        print(_node)
        raise RuntimeError('Failed to find state for node')

    x = _node['x']
    molWmix = float(x[0]) * float(molWNH3) + float(x[1]) * float(molWH2O)

    try:
        propl = rp.flsh('ph', _node['p'], _node['h'] - 25*molWmix*depth, _node['x'])
        propr = rp.flsh('ph', _node['p'], _node['h'] + 25*molWmix*depth, _node['x'])
        t = (propl['t'] + propr['t']) / 2

        prop = rp.flsh('tp', t, _node['p'], _node['x'])
    except rp.RefpropError as e:
        _depth = depth + 1
        prop = _PHworkaround(_node, depth = _depth)

    return prop

def refpropState(node):
    '''
    If the state is to be found from refprop, use this method
    '''

    mode = ''
    in1 = 0
    in2 = 0

    # Convert input to refprop notation:
    _node = toRefprop(node)

    # Figure out which inputs to use (sorted by priority)
    if 'p' in _node and 'h' in _node:
        mode = 'ph'
        in1 = _node['p']
        in2 = _node['h']
    elif 'p' in _node and 'e' in _node:
        mode = 'pe'
        in1 = _node['p']
        in2 = _node['e']
    elif 'p' in _node and 's' in _node:
        mode = 'ps'
        in1 = _node['p']
        in2 = _node['s']
    elif 'p' in _node and 'q' in _node:
        mode = 'pq'
        in1 = _node['p']
        in2 = _node['q']
    elif 't' in _node and 'p' in _node:
        mode = 'tp'
        in1 = _node['t']
        in2 = _node['p']
    elif 't' in _node and 'q' in _node:
        mode = 'tq'
        in1 = _node['t']
        in2 = _node['q']
    else:
        print(_node)
        raise InputError('state','Missing inputs for above node')

    try:
        # Calculate
        prop = rp.flsh(mode, in1, in2, _node['x'])
    except rp.RefpropError as e:
        # Normal flsh failed, try flsh2 as well
        try:
            if mode is 'ph':
                # Assume something in the 2-phase region. temperature changes very little with enthalpy
                prop = _PHworkaround(_node)
                print('Workaround for iteration, enthalpy difference {:.2e} K'.format(in2 - _node['h']))
            elif mode is 'tp':
                propl = rp.flsh('tp', _node['t'] - 0.1, _node['p'], _node['x'])
                propr = rp.flsh('tp', _node['t'] + 0.1, _node['p'], _node['x'])
                h = (propl['h'] + propr['h']) / 2

                prop = rp.flsh('ph', _node['p'], h, _node['x'])
                print('Workaround for iteration, temperature difference {:.2e} J/mol*K'.format(in1 - _node['t']))
            else:
                print(node)
                print(mode)
                print(in1)
                print(in2)
                raise(e)

        except Exception as e:
            print(node)
            raise(e)

    # Convert back from refprop notation, then update node
    node.update(fromRefprop(prop))

    return node

def cpBasedState(node):
    '''
    This does not consider:
    p,cv,q,y,yvap,yliq
    But it considers:
    e,h,s,t,cp
    Reference point is: h = 0 at t=273
    '''

    # Make sure object has often-requested properties defined
    if not 'p' in node:
        node['p'] = 0

    if not 'q' in node:
        node['q'] = 0

    if not 'y' in node:
        node['y'] = 0

    # Calculation
    if 'h' in node:
        node['t'] = node['h'] / node['cp']

    elif 't' in node:
        node['h'] = node['cp'] * node['t']

    return node
<commit_msg>Also calculate entropy for self-defined media<commit_after>from decimal import Decimal

from dna.vendor import refprop as rp

class Error(Exception):
    """Base class for exceptions in this module."""
    pass

class InputError(Error):
    """Exception raised for errors in the input.

    Attributes:
        expr -- input expression in which the error occurred
        msg  -- explanation of the error
    """

    def __init__(self, expr, msg):
        self.expr = expr
        self.msg = msg

rp.setup('def', 'ammonia', 'water')
rp.setref(hrf='def',ixflag=1)
molWNH3 = rp.info(1)['wmm']
molWH2O = rp.info(2)['wmm']

def state(node):

    use_rp = True

    if 'media' in node:
        usermedia = {
            'hitecxl': {'cp': 1.447, 'tmin': 130, 'tmax': 490},
            'hitec': {'cp': 1.5617, 'tmin': 180, 'tmax': 560}
        }

        if node['media'] in usermedia:
            node['cp'] = usermedia[node['media']]['cp']
            node['tmin'] = usermedia[node['media']]['tmin']
            node['tmax'] = usermedia[node['media']]['tmax']
            use_rp = False

    if use_rp:
        return refpropState(node)
    else:
        return cpBasedState(node)

def molarToMass(x):
    mass_nh3 = float(x)*molWNH3;
    mass_h2o = (1-float(x))*molWH2O;

    mass_total = mass_nh3 + mass_h2o

    return [mass_nh3/mass_total, 1-mass_nh3/mass_total]

def massToMolar(y):
    m_nh3 = float(y)/molWNH3
    m_h2o = (1-float(y))/molWH2O

    mt = m_nh3 + m_h2o

    return [m_nh3/mt, 1-m_nh3/mt]

def toRefprop(node):
    prop = {}

    x = massToMolar(node['y'])

    molWmix = float(x[0]) * float(molWNH3) + float(x[1]) * float(molWH2O) # g/mol

    if 'mdot' in node:
        prop['mdot'] = node['mdot']

    if 'p' in node:
        prop['p'] = node['p']*100 # hPa > kPa

    if 'e' in node:
        prop['e'] = node['e']*molWmix # kJ/kg > J/mol

    if 'h' in node:
        prop['h'] = node['h']*molWmix # kJ/kg*K > J/mol*K

    if 'D' in node:
        prop['D'] = node['D']/molWmix*1000 # kg/L > mol/m3

    if 's' in node:
        prop['s'] = node['s']*molWmix # kJ/kg*K > J/mol*K

    if 't' in node:
        prop['t'] = node['t'] + 273.15 # C > K

    if 'cv' in node:
        prop['cv'] = node['cv']*molWmix # kJ/kg*K > J/mol*K

    if 'cp' in node:
        prop['cp'] = node['cp']*molWmix # kJ/kg*K > J/mol*K

    if 'q' in node:
        prop['q'] = massToMolar(node['q'])[0] # kg/kg > mol/mol

    if 'y' in node:
        prop['x'] = massToMolar(node['y']) # kg/kg > mol/mol

    if 'yvap' in node:
        prop['xvap'] = massToMolar(node['yvap']) # kg/kg > mol/mol

    if 'yliq' in node:
        prop['xliq'] = massToMolar(node['yliq']) # kg/kg > mol/mol

    return prop

def fromRefprop(prop):
    node = {}

    x = prop['x']

    molWmix = float(x[0]) * float(molWNH3) + float(x[1]) * float(molWH2O)

    if 'mdot' in prop:
        node['mdot'] = prop['mdot']

    # Convert mol to kg, K to C, kPa to hPa
    if 'p' in prop:
        node['p'] = prop['p']/100 # kPa > hPa

    if 'e' in prop:
        node['e'] = prop['e']/molWmix # J/mol > kJ/kg

    if 'h' in prop:
        node['h'] = prop['h']/molWmix # J/mol > kJ/kg

    if 'D' in prop:
        node['D'] = prop['D']*molWmix/1000 # mol/L > kg/m3

    if 's' in prop:
        node['s'] = prop['s']/molWmix # J/mol*K > kJ/kg*K

    if 't' in prop:
        node['t'] = prop['t'] - 273.15 # K > C

    if 'cv' in prop:
        node['cv'] = prop['cv']/molWmix # J/mol*K > kJ/kg*K

    if 'cp' in prop:
        node['cp'] = prop['cp']/molWmix # J/mol*K > kJ/kg*K

    if 'q' in prop:
        node['q'] = molarToMass(prop['q'])[0] # mol/mol > kg/kg

    if 'x' in prop:
        node['y'] = molarToMass(prop['x'][0])[0] # mol/mol > kg/kg

    if 'xvap' in prop:
        node['yvap'] = molarToMass(prop['xvap'][0])[0] # mol/mol > kg/kg

    if 'xliq' in prop:
        node['yliq'] = molarToMass(prop['xliq'][0])[0] # mol/mol > kg/kg

    return node

def _PHworkaround(_node, depth = 1):

    if depth > 5:
        print(_node)
        raise RuntimeError('Failed to find state for node')

    x = _node['x']
    molWmix = float(x[0]) * float(molWNH3) + float(x[1]) * float(molWH2O)

    try:
        propl = rp.flsh('ph', _node['p'], _node['h'] - 25*molWmix*depth, _node['x'])
        propr = rp.flsh('ph', _node['p'], _node['h'] + 25*molWmix*depth, _node['x'])
        t = (propl['t'] + propr['t']) / 2

        prop = rp.flsh('tp', t, _node['p'], _node['x'])
    except rp.RefpropError as e:
        _depth = depth + 1
        prop = _PHworkaround(_node, depth = _depth)

    return prop

def refpropState(node):
    '''
    If the state is to be found from refprop, use this method
    '''

    mode = ''
    in1 = 0
    in2 = 0

    # Convert input to refprop notation:
    _node = toRefprop(node)

    # Figure out which inputs to use (sorted by priority)
    if 'p' in _node and 'h' in _node:
        mode = 'ph'
        in1 = _node['p']
        in2 = _node['h']
    elif 'p' in _node and 'e' in _node:
        mode = 'pe'
        in1 = _node['p']
        in2 = _node['e']
    elif 'p' in _node and 's' in _node:
        mode = 'ps'
        in1 = _node['p']
        in2 = _node['s']
    elif 'p' in _node and 'q' in _node:
        mode = 'pq'
        in1 = _node['p']
        in2 = _node['q']
    elif 't' in _node and 'p' in _node:
        mode = 'tp'
        in1 = _node['t']
        in2 = _node['p']
    elif 't' in _node and 'q' in _node:
        mode = 'tq'
        in1 = _node['t']
        in2 = _node['q']
    else:
        print(_node)
        raise InputError('state','Missing inputs for above node')

    try:
        # Calculate
        prop = rp.flsh(mode, in1, in2, _node['x'])
    except rp.RefpropError as e:
        # Normal flsh failed, try flsh2 as well
        try:
            if mode is 'ph':
                # Assume something in the 2-phase region. temperature changes very little with enthalpy
                prop = _PHworkaround(_node)
                print('Workaround for iteration, enthalpy difference {:.2e} K'.format(in2 - _node['h']))
            elif mode is 'tp':
                propl = rp.flsh('tp', _node['t'] - 0.1, _node['p'], _node['x'])
                propr = rp.flsh('tp', _node['t'] + 0.1, _node['p'], _node['x'])
                h = (propl['h'] + propr['h']) / 2

                prop = rp.flsh('ph', _node['p'], h, _node['x'])
                print('Workaround for iteration, temperature difference {:.2e} J/mol*K'.format(in1 - _node['t']))
            else:
                print(node)
                print(mode)
                print(in1)
                print(in2)
                raise(e)

        except Exception as e:
            print(node)
            raise(e)

    # Convert back from refprop notation, then update node
    node.update(fromRefprop(prop))

    return node

def cpBasedState(node):
    '''
    This does not consider:
    p,cv,q,y,yvap,yliq
    But it considers:
    e,h,s,t,cp
    Reference point is: h = 0 at t=273
    '''

    # Make sure object has often-requested properties defined
    if not 'p' in node:
        node['p'] = 0

    if not 'q' in node:
        node['q'] = 0

    if not 'y' in node:
        node['y'] = 0

    # Calculation
    if 'h' in node:
        node['t'] = node['h'] / node['cp']
        node['s'] = node['h'] / (node['t'] + 273.15)

    elif 't' in node:
        node['h'] = node['cp'] * node['t']
        node['s'] = node['h'] / (node['t'] + 273.15)
    elif 's' in node:
        node['t'] = 273.15 / ((node['cp'] / node['s']) - 1)
        node['h'] = node['cp'] * node['t']

    return node
<|endoftext|>


<commit_before>from __future__ import unicode_literals

import pygst
pygst.require('0.10')
import gst
import gobject

import ConfigParser as configparser
import io

try:
    import xml.etree.cElementTree as elementtree
except ImportError:
    import xml.etree.ElementTree as elementtree

# TODO: make detect_FOO_header reusable in general mopidy code.
# i.e. give it just a "peek" like function.
def detect_m3u_header(typefind):
    return typefind.peek(0, 8) == b'#EXTM3U\n'

def detect_pls_header(typefind):
    return typefind.peek(0, 11) == b'[playlist]\n'

def detect_xspf_header(typefind):
    # Get more data than the 90 needed for header in case spacing is funny.
    data = io.BytesIO(typefind.peek(0, 150))
    try:
        for event, element in elementtree.iterparse(data, events=('start',)):
            return element.tag.lower() == '{http://xspf.org/ns/0/}playlist'
    except elementtree.ParseError:
        pass
    return False

def detect_asx_header(typefind):
    data = io.BytesIO(typefind.peek(0, 50))
    try:
        for event, element in elementtree.iterparse(data, events=('start',)):
            return element.tag.lower() == 'asx'
    except elementtree.ParseError:
        pass
    return False

def parse_m3u(data):
    # TODO: convert non URIs to file URIs.
    for line in data.readlines():
        if not line.startswith('#') and line.strip():
            yield line

def parse_pls(data):
    # TODO: error handling of bad playlists.
    # TODO: convert non URIs to file URIs.
    cp = configparser.RawConfigParser()
    cp.readfp(data)
    for i in xrange(1, cp.getint('playlist', 'numberofentries')):
        yield cp.get('playlist', 'file%d' % i)

def parse_xspf(data):
    # TODO: handle parser errors
    for event, element in elementtree.iterparse(data):
        element.tag = element.tag.lower()  # normalize

    ns = 'http://xspf.org/ns/0/'
    for track in element.iterfind('{%s}tracklist/{%s}track' % (ns, ns)):
        yield track.findtext('{%s}location' % ns)

def parse_asx(data):
    # TODO: handle parser errors
    for event, element in elementtree.iterparse(data):
        element.tag = element.tag.lower()  # normalize

    for ref in element.findall('entry/ref'):
        yield ref.get('href', '').strip()

def parse_urilist(data):
    for line in data.readlines():
        if not line.startswith('#') and gst.uri_is_valid(line.strip()):
            yield line

def playlist_typefinder(typefind, func, caps):
    if func(typefind):
        typefind.suggest(gst.TYPE_FIND_MAXIMUM, caps)

def register_typefind(mimetype, func, extensions):
    caps = gst.caps_from_string(mimetype)
    gst.type_find_register(mimetype, gst.RANK_PRIMARY, playlist_typefinder,
                           extensions, caps, func, caps)

def register_typefinders():
    register_typefind('audio/x-mpegurl', detect_m3u_header, [b'm3u', b'm3u8'])
    register_typefind('audio/x-scpls', detect_pls_header, [b'pls'])
    register_typefind('application/xspf+xml', detect_xspf_header, [b'xspf'])
    # NOTE: seems we can't use video/x-ms-asf which is the correct mime for asx
    # as it is shared with asf for streaming videos :/
    register_typefind('audio/x-ms-asx', detect_asx_header, [b'asx'])

class BasePlaylistElement(gst.Bin):
    """Base class for creating GStreamer elements for playlist support.

    This element performs the following steps:

    1. Initializes src and sink pads for the element.
    2. Collects data from the sink until EOS is reached.
    3. Passes the collected data to :meth:`convert` to get a list of URIs.
    4. Passes the list of URIs to :meth:`handle`, default handling is to pass
       the URIs to the src element as a uri-list.
    5. If handle returned true, the EOS consumed and nothing more happens, if
       it is not consumed it flows on to the next element downstream, which is
       likely our uri-list consumer which needs the EOS to know we are done
       sending URIs.
    """

    sinkpad_template = None
    """GStreamer pad template to use for sink, must be overriden."""

    srcpad_template = None
    """GStreamer pad template to use for src, must be overriden."""

    ghost_srcpad = False
    """Indicates if src pad should be ghosted or not."""

    def __init__(self):
        """Sets up src and sink pads plus behaviour."""
        super(BasePlaylistElement, self).__init__()
        self._data = io.BytesIO()
        self._done = False

        self.sinkpad = gst.Pad(self.sinkpad_template)
        self.sinkpad.set_chain_function(self._chain)
        self.sinkpad.set_event_function(self._event)
        self.add_pad(self.sinkpad)

        if self.ghost_srcpad:
            self.srcpad = gst.ghost_pad_new_notarget('src', gst.PAD_SRC)
        else:
            self.srcpad = gst.Pad(self.srcpad_template)
        self.add_pad(self.srcpad)

    def convert(self, data):
        """Convert the data we have colleted to URIs.

        :param data: collected data buffer
        :type data: :class:`io.BytesIO`
        :returns: iterable or generator of URIs
        """
        raise NotImplementedError

    def handle(self, uris):
        """Do something useful with the URIs.

        :param uris: list of URIs
        :type uris: :type:`list`
        :returns: boolean indicating if EOS should be consumed
        """
        # TODO: handle unicode uris which we can get out of elementtree
        self.srcpad.push(gst.Buffer('\n'.join(uris)))
        return False

    def _chain(self, pad, buf):
        if not self._done:
            self._data.write(buf.data)
            return gst.FLOW_OK
        return gst.FLOW_EOS

    def _event(self, pad, event):
        if event.type == gst.EVENT_NEWSEGMENT:
            return True

        if event.type == gst.EVENT_EOS:
            self._done = True
            self._data.seek(0)
            if self.handle(list(self.convert(self._data))):
                return True

        # Ensure we handle remaining events in a sane way.
        return pad.event_default(event)

class M3uDecoder(BasePlaylistElement):
    __gstdetails__ = ('M3U Decoder',
                      'Decoder',
                      'Convert .m3u to text/uri-list',
                      'Mopidy')

    sinkpad_template = gst.PadTemplate('sink',
        gst.PAD_SINK,
        gst.PAD_ALWAYS,
        gst.caps_from_string('audio/x-mpegurl'))

    srcpad_template = gst.PadTemplate('src',
        gst.PAD_SRC,
        gst.PAD_ALWAYS,
        gst.caps_from_string('text/uri-list'))

    __gsttemplates__ = (sinkpad_template, srcpad_template)

    def convert(self, data):
        return parse_m3u(data)

class PlsDecoder(BasePlaylistElement):
    __gstdetails__ = ('PLS Decoder',
                      'Decoder',
                      'Convert .pls to text/uri-list',
                      'Mopidy')

    sinkpad_template = gst.PadTemplate('sink',
        gst.PAD_SINK,
        gst.PAD_ALWAYS,
        gst.caps_from_string('audio/x-scpls'))

    srcpad_template = gst.PadTemplate('src',
        gst.PAD_SRC,
        gst.PAD_ALWAYS,
        gst.caps_from_string('text/uri-list'))

    __gsttemplates__ = (sinkpad_template, srcpad_template)

    def convert(self, data):
        return parse_pls(data)

class XspfDecoder(BasePlaylistElement):
    __gstdetails__ = ('XSPF Decoder',
                      'Decoder',
                      'Convert .pls to text/uri-list',
                      'Mopidy')

    sinkpad_template = gst.PadTemplate('sink',
        gst.PAD_SINK,
        gst.PAD_ALWAYS,
        gst.caps_from_string('application/xspf+xml'))

    srcpad_template = gst.PadTemplate('src',
        gst.PAD_SRC,
        gst.PAD_ALWAYS,
        gst.caps_from_string('text/uri-list'))

    __gsttemplates__ = (sinkpad_template, srcpad_template)

    def convert(self, data):
        return parse_xspf(data)

class AsxDecoder(BasePlaylistElement):
    __gstdetails__ = ('ASX Decoder',
                      'Decoder',
                      'Convert .asx to text/uri-list',
                      'Mopidy')

    sinkpad_template = gst.PadTemplate('sink',
        gst.PAD_SINK,
        gst.PAD_ALWAYS,
        gst.caps_from_string('audio/x-ms-asx'))

    srcpad_template = gst.PadTemplate('src',
        gst.PAD_SRC,
        gst.PAD_ALWAYS,
        gst.caps_from_string('text/uri-list'))

    __gsttemplates__ = (sinkpad_template, srcpad_template)

    def convert(self, data):
        return parse_asx(data)

class UriListElement(BasePlaylistElement):
    __gstdetails__ = ('URIListDemuxer',
                      'Demuxer',
                      'Convert a text/uri-list to a stream',
                      'Mopidy')

    sinkpad_template = gst.PadTemplate('sink',
        gst.PAD_SINK,
        gst.PAD_ALWAYS,
        gst.caps_from_string('text/uri-list'))

    srcpad_template = gst.PadTemplate('src',
        gst.PAD_SRC,
        gst.PAD_ALWAYS,
        gst.caps_new_any())

    ghost_srcpad = True  # We need to hook this up to our internal decodebin

    __gsttemplates__ = (sinkpad_template, srcpad_template)

    def __init__(self):
        super(UriListElement, self).__init__()
        self.uridecodebin = gst.element_factory_make('uridecodebin')
        self.uridecodebin.connect('pad-added', self.pad_added)
        # Limit to anycaps so we get a single stream out, letting other
        # elements downstream figure out actual muxing
        self.uridecodebin.set_property('caps', gst.caps_new_any())

    def pad_added(self, src, pad):
        self.srcpad.set_target(pad)

    def handle(self, uris):
        # TODO: hookup about to finish and errors to rest of URIs so we
        # round robin, only giving up once all have been tried.
        # TODO: uris could be empty.
        self.add(self.uridecodebin)
        self.uridecodebin.set_state(gst.STATE_READY)
        self.uridecodebin.set_property('uri', uris[0])
        self.uridecodebin.sync_state_with_parent()
        return True  # Make sure we consume the EOS that triggered us.

    def convert(self, data):
        return parse_urilist(data)

class IcySrc(gst.Bin, gst.URIHandler):
    __gstdetails__ = ('IcySrc',
                      'Src',
                      'HTTP src wrapper for icy:// support.',
                      'Mopidy')

    srcpad_template = gst.PadTemplate('src',
        gst.PAD_SRC,
        gst.PAD_ALWAYS,
        gst.caps_new_any())

    __gsttemplates__ = (srcpad_template,)

    def __init__(self):
        super(IcySrc, self).__init__()
        self._httpsrc = gst.element_make_from_uri(gst.URI_SRC, 'http://')
        try:
            self._httpsrc.set_property('iradio-mode', True)
        except TypeError:
            pass
        self.add(self._httpsrc)

        self._srcpad = gst.GhostPad('src', self._httpsrc.get_pad('src'))
        self.add_pad(self._srcpad)

    @classmethod
    def do_get_type_full(cls):
        return gst.URI_SRC

    @classmethod
    def do_get_protocols_full(cls):
        return [b'icy', b'icyx']

    def do_set_uri(self, uri):
        if uri.startswith('icy://'):
            return self._httpsrc.set_uri(b'http://' + uri[len('icy://'):])
        elif uri.startswith('icyx://'):
            return self._httpsrc.set_uri(b'https://' + uri[len('icyx://'):])
        else:
            return False

    def do_get_uri(self):
        uri = self._httpsrc.get_uri()
        if uri.startswith('http://'):
            return b'icy://' + uri[len('http://'):]
        else:
            return b'icyx://' + uri[len('https://'):]

def register_element(element_class):
    gobject.type_register(element_class)
    gst.element_register(
        element_class, element_class.__name__.lower(), gst.RANK_MARGINAL)

def register_elements():
    register_element(M3uDecoder)
    register_element(PlsDecoder)
    register_element(XspfDecoder)
    register_element(AsxDecoder)
    register_element(UriListElement)
    if not gst.element_make_from_uri(gst.URI_SRC, 'icy://'):
        register_element(IcySrc)
<commit_msg>audio: Post an error if an urilist expands to another urilist<commit_after>from __future__ import unicode_literals

import pygst
pygst.require('0.10')
import gst
import gobject

import ConfigParser as configparser
import io

try:
    import xml.etree.cElementTree as elementtree
except ImportError:
    import xml.etree.ElementTree as elementtree

# TODO: make detect_FOO_header reusable in general mopidy code.
# i.e. give it just a "peek" like function.
def detect_m3u_header(typefind):
    return typefind.peek(0, 8) == b'#EXTM3U\n'

def detect_pls_header(typefind):
    return typefind.peek(0, 11) == b'[playlist]\n'

def detect_xspf_header(typefind):
    # Get more data than the 90 needed for header in case spacing is funny.
    data = io.BytesIO(typefind.peek(0, 150))
    try:
        for event, element in elementtree.iterparse(data, events=('start',)):
            return element.tag.lower() == '{http://xspf.org/ns/0/}playlist'
    except elementtree.ParseError:
        pass
    return False

def detect_asx_header(typefind):
    data = io.BytesIO(typefind.peek(0, 50))
    try:
        for event, element in elementtree.iterparse(data, events=('start',)):
            return element.tag.lower() == 'asx'
    except elementtree.ParseError:
        pass
    return False

def parse_m3u(data):
    # TODO: convert non URIs to file URIs.
    for line in data.readlines():
        if not line.startswith('#') and line.strip():
            yield line

def parse_pls(data):
    # TODO: error handling of bad playlists.
    # TODO: convert non URIs to file URIs.
    cp = configparser.RawConfigParser()
    cp.readfp(data)
    for i in xrange(1, cp.getint('playlist', 'numberofentries')):
        yield cp.get('playlist', 'file%d' % i)

def parse_xspf(data):
    # TODO: handle parser errors
    for event, element in elementtree.iterparse(data):
        element.tag = element.tag.lower()  # normalize

    ns = 'http://xspf.org/ns/0/'
    for track in element.iterfind('{%s}tracklist/{%s}track' % (ns, ns)):
        yield track.findtext('{%s}location' % ns)

def parse_asx(data):
    # TODO: handle parser errors
    for event, element in elementtree.iterparse(data):
        element.tag = element.tag.lower()  # normalize

    for ref in element.findall('entry/ref'):
        yield ref.get('href', '').strip()

def parse_urilist(data):
    for line in data.readlines():
        if not line.startswith('#') and gst.uri_is_valid(line.strip()):
            yield line

def playlist_typefinder(typefind, func, caps):
    if func(typefind):
        typefind.suggest(gst.TYPE_FIND_MAXIMUM, caps)

def register_typefind(mimetype, func, extensions):
    caps = gst.caps_from_string(mimetype)
    gst.type_find_register(mimetype, gst.RANK_PRIMARY, playlist_typefinder,
                           extensions, caps, func, caps)

def register_typefinders():
    register_typefind('audio/x-mpegurl', detect_m3u_header, [b'm3u', b'm3u8'])
    register_typefind('audio/x-scpls', detect_pls_header, [b'pls'])
    register_typefind('application/xspf+xml', detect_xspf_header, [b'xspf'])
    # NOTE: seems we can't use video/x-ms-asf which is the correct mime for asx
    # as it is shared with asf for streaming videos :/
    register_typefind('audio/x-ms-asx', detect_asx_header, [b'asx'])

class BasePlaylistElement(gst.Bin):
    """Base class for creating GStreamer elements for playlist support.

    This element performs the following steps:

    1. Initializes src and sink pads for the element.
    2. Collects data from the sink until EOS is reached.
    3. Passes the collected data to :meth:`convert` to get a list of URIs.
    4. Passes the list of URIs to :meth:`handle`, default handling is to pass
       the URIs to the src element as a uri-list.
    5. If handle returned true, the EOS consumed and nothing more happens, if
       it is not consumed it flows on to the next element downstream, which is
       likely our uri-list consumer which needs the EOS to know we are done
       sending URIs.
    """

    sinkpad_template = None
    """GStreamer pad template to use for sink, must be overriden."""

    srcpad_template = None
    """GStreamer pad template to use for src, must be overriden."""

    ghost_srcpad = False
    """Indicates if src pad should be ghosted or not."""

    def __init__(self):
        """Sets up src and sink pads plus behaviour."""
        super(BasePlaylistElement, self).__init__()
        self._data = io.BytesIO()
        self._done = False

        self.sinkpad = gst.Pad(self.sinkpad_template)
        self.sinkpad.set_chain_function(self._chain)
        self.sinkpad.set_event_function(self._event)
        self.add_pad(self.sinkpad)

        if self.ghost_srcpad:
            self.srcpad = gst.ghost_pad_new_notarget('src', gst.PAD_SRC)
        else:
            self.srcpad = gst.Pad(self.srcpad_template)
        self.add_pad(self.srcpad)

    def convert(self, data):
        """Convert the data we have colleted to URIs.

        :param data: collected data buffer
        :type data: :class:`io.BytesIO`
        :returns: iterable or generator of URIs
        """
        raise NotImplementedError

    def handle(self, uris):
        """Do something useful with the URIs.

        :param uris: list of URIs
        :type uris: :type:`list`
        :returns: boolean indicating if EOS should be consumed
        """
        # TODO: handle unicode uris which we can get out of elementtree
        self.srcpad.push(gst.Buffer('\n'.join(uris)))
        return False

    def _chain(self, pad, buf):
        if not self._done:
            self._data.write(buf.data)
            return gst.FLOW_OK
        return gst.FLOW_EOS

    def _event(self, pad, event):
        if event.type == gst.EVENT_NEWSEGMENT:
            return True

        if event.type == gst.EVENT_EOS:
            self._done = True
            self._data.seek(0)
            if self.handle(list(self.convert(self._data))):
                return True

        # Ensure we handle remaining events in a sane way.
        return pad.event_default(event)

class M3uDecoder(BasePlaylistElement):
    __gstdetails__ = ('M3U Decoder',
                      'Decoder',
                      'Convert .m3u to text/uri-list',
                      'Mopidy')

    sinkpad_template = gst.PadTemplate('sink',
        gst.PAD_SINK,
        gst.PAD_ALWAYS,
        gst.caps_from_string('audio/x-mpegurl'))

    srcpad_template = gst.PadTemplate('src',
        gst.PAD_SRC,
        gst.PAD_ALWAYS,
        gst.caps_from_string('text/uri-list'))

    __gsttemplates__ = (sinkpad_template, srcpad_template)

    def convert(self, data):
        return parse_m3u(data)

class PlsDecoder(BasePlaylistElement):
    __gstdetails__ = ('PLS Decoder',
                      'Decoder',
                      'Convert .pls to text/uri-list',
                      'Mopidy')

    sinkpad_template = gst.PadTemplate('sink',
        gst.PAD_SINK,
        gst.PAD_ALWAYS,
        gst.caps_from_string('audio/x-scpls'))

    srcpad_template = gst.PadTemplate('src',
        gst.PAD_SRC,
        gst.PAD_ALWAYS,
        gst.caps_from_string('text/uri-list'))

    __gsttemplates__ = (sinkpad_template, srcpad_template)

    def convert(self, data):
        return parse_pls(data)

class XspfDecoder(BasePlaylistElement):
    __gstdetails__ = ('XSPF Decoder',
                      'Decoder',
                      'Convert .pls to text/uri-list',
                      'Mopidy')

    sinkpad_template = gst.PadTemplate('sink',
        gst.PAD_SINK,
        gst.PAD_ALWAYS,
        gst.caps_from_string('application/xspf+xml'))

    srcpad_template = gst.PadTemplate('src',
        gst.PAD_SRC,
        gst.PAD_ALWAYS,
        gst.caps_from_string('text/uri-list'))

    __gsttemplates__ = (sinkpad_template, srcpad_template)

    def convert(self, data):
        return parse_xspf(data)

class AsxDecoder(BasePlaylistElement):
    __gstdetails__ = ('ASX Decoder',
                      'Decoder',
                      'Convert .asx to text/uri-list',
                      'Mopidy')

    sinkpad_template = gst.PadTemplate('sink',
        gst.PAD_SINK,
        gst.PAD_ALWAYS,
        gst.caps_from_string('audio/x-ms-asx'))

    srcpad_template = gst.PadTemplate('src',
        gst.PAD_SRC,
        gst.PAD_ALWAYS,
        gst.caps_from_string('text/uri-list'))

    __gsttemplates__ = (sinkpad_template, srcpad_template)

    def convert(self, data):
        return parse_asx(data)

class UriListElement(BasePlaylistElement):
    __gstdetails__ = ('URIListDemuxer',
                      'Demuxer',
                      'Convert a text/uri-list to a stream',
                      'Mopidy')

    sinkpad_template = gst.PadTemplate('sink',
        gst.PAD_SINK,
        gst.PAD_ALWAYS,
        gst.caps_from_string('text/uri-list'))

    srcpad_template = gst.PadTemplate('src',
        gst.PAD_SRC,
        gst.PAD_ALWAYS,
        gst.caps_new_any())

    ghost_srcpad = True  # We need to hook this up to our internal decodebin

    __gsttemplates__ = (sinkpad_template, srcpad_template)

    def __init__(self):
        super(UriListElement, self).__init__()
        self.uridecodebin = gst.element_factory_make('uridecodebin')
        self.uridecodebin.connect('pad-added', self.pad_added)
        # Limit to anycaps so we get a single stream out, letting other
        # elements downstream figure out actual muxing
        self.uridecodebin.set_property('caps', gst.caps_new_any())

    def pad_added(self, src, pad):
        self.srcpad.set_target(pad)
        pad.add_event_probe(self.pad_event)

    def pad_event(self, pad, event):
        if event.has_name('urilist-played'):
            error = gst.GError(gst.RESOURCE_ERROR, gst.RESOURCE_ERROR_FAILED,
                               b'Nested playlists not supported.')
            message = gst.message_new_error(self, error, b'Playlists pointing to other playlists is not supported')
            self.post_message(message)
        return True

    def handle(self, uris):
        struct = gst.Structure('urilist-played')
        event = gst.event_new_custom(gst.EVENT_CUSTOM_UPSTREAM, struct)
        self.sinkpad.push_event(event)

        # TODO: hookup about to finish and errors to rest of URIs so we
        # round robin, only giving up once all have been tried.
        # TODO: uris could be empty.
        self.add(self.uridecodebin)
        self.uridecodebin.set_state(gst.STATE_READY)
        self.uridecodebin.set_property('uri', uris[0])
        self.uridecodebin.sync_state_with_parent()
        return True  # Make sure we consume the EOS that triggered us.

    def convert(self, data):
        return parse_urilist(data)

class IcySrc(gst.Bin, gst.URIHandler):
    __gstdetails__ = ('IcySrc',
                      'Src',
                      'HTTP src wrapper for icy:// support.',
                      'Mopidy')

    srcpad_template = gst.PadTemplate('src',
        gst.PAD_SRC,
        gst.PAD_ALWAYS,
        gst.caps_new_any())

    __gsttemplates__ = (srcpad_template,)

    def __init__(self):
        super(IcySrc, self).__init__()
        self._httpsrc = gst.element_make_from_uri(gst.URI_SRC, 'http://')
        try:
            self._httpsrc.set_property('iradio-mode', True)
        except TypeError:
            pass
        self.add(self._httpsrc)

        self._srcpad = gst.GhostPad('src', self._httpsrc.get_pad('src'))
        self.add_pad(self._srcpad)

    @classmethod
    def do_get_type_full(cls):
        return gst.URI_SRC

    @classmethod
    def do_get_protocols_full(cls):
        return [b'icy', b'icyx']

    def do_set_uri(self, uri):
        if uri.startswith('icy://'):
            return self._httpsrc.set_uri(b'http://' + uri[len('icy://'):])
        elif uri.startswith('icyx://'):
            return self._httpsrc.set_uri(b'https://' + uri[len('icyx://'):])
        else:
            return False

    def do_get_uri(self):
        uri = self._httpsrc.get_uri()
        if uri.startswith('http://'):
            return b'icy://' + uri[len('http://'):]
        else:
            return b'icyx://' + uri[len('https://'):]

def register_element(element_class):
    gobject.type_register(element_class)
    gst.element_register(
        element_class, element_class.__name__.lower(), gst.RANK_MARGINAL)

def register_elements():
    register_element(M3uDecoder)
    register_element(PlsDecoder)
    register_element(XspfDecoder)
    register_element(AsxDecoder)
    register_element(UriListElement)
    if not gst.element_make_from_uri(gst.URI_SRC, 'icy://'):
        register_element(IcySrc)
<|endoftext|>


<commit_before>"""
  storm SiteMover

  :author: Mario Lassnig <mario.lassnig@cern.ch>, 2016
"""

from .base import BaseSiteMover

from TimerCommand import TimerCommand
from PilotErrors import PilotException

from rucio.client import ReplicaClient

from datetime import datetime
from xml.dom import minidom

import os, shutil

class stormSiteMover(BaseSiteMover):
    """
    SiteMover that queries HTTP etag for physcial storage location, then
    symlinks for stage in, and copies for stage out
    """

    name = 'storm'
    schemes = ['file']
    version = '20161125.005'

    require_replicas = False       ## quick hack to avoid query Rucio to resolve input replicas

    def __init__(self, *args, **kwargs):
        super(stormSiteMover, self).__init__(*args, **kwargs)
        self.log('storm sitemover version: %s' % self.version)

    def resolve_replica(self, fspec, protocol, ddm=None):
        """
        Overridden method -- unused
        """

        return {'ddmendpoint': fspec.replicas[0][0] if fspec.replicas else None,
                'surl': None,
                'pfn': fspec.lfn}

    def stageIn(self, source, destination, fspec):
        """
        Query HTTP for etag, then symlink to the pilot working directory.

        :param source:      original file location
        :param destination: where to create the link
        :param fspec:       dictionary containing destination replicas, scope, lfn
        :return:            destination file details (checksumtype, checksum, size)
        """

        self.log('source: %s' % str(source))
        self.log('destination: %s' % str(destination))
        self.log('fspec: %s' % str(fspec))
        self.log('fspec.scope: %s' % str(fspec.scope))
        self.log('fspec.lfn: %s' % str(fspec.lfn))
        self.log('fspec.ddmendpoint: %s' % str(fspec.ddmendpoint))

        # figure out the HTTP SURL from Rucio
        rc = ReplicaClient()
        http_surl_reps = [r for r in rc.list_replicas(dids=[{'scope': fspec.scope,
                                                             'name': fspec.lfn}],
                                                      schemes=['https'],
                                                      rse_expression=fspec.ddmendpoint)]
        self.log('http_surl_reps: %s' % http_surl_reps)

        http_surl = http_surl_reps[0]['rses'][fspec.ddmendpoint][0].rsplit('_-')[0]
        self.log('http_surl: %s' % http_surl)

        # retrieve the TURL from the webdav etag
        cmd = 'davix-http --capath /cvmfs/atlas.cern.ch/repo/ATLASLocalRootBase/etc/grid-security-emi/certificates --cert $X509_USER_PROXY -X PROPFIND %s' % http_surl
        self.log('ETAG retrieval: %s' % cmd)
        try:
            timer = TimerCommand(cmd)
            rcode, output = timer.run(timeout=10)
        except Exception, e:
            self.log('FATAL: could not retrieve STORM WebDAV ETag: %s' % e)
            raise PilotException('Could not retrieve STORM WebDAV ETag: %s' % e)
        p_output = minidom.parseString(output)

        # we need to strip off the quotation marks and the <timestamp> from the etag
        # but since we can have multiple underscores, we have to rely on the uniqueness
        # of the full LFN to make the split
        target = p_output.getElementsByTagName('d:getetag')[0].childNodes[0].nodeValue.replace('"', '')
        self.log('Symlink before: %s' % target)
        target = target.split(fspec.lfn)[0]+fspec.lfn
        self.log('Symlink after : %s' % target)

        # make the symlink
        try:
            self.log('Making symlink from %s to %s' % (target, destination))
            os.symlink(target, destination)
        except Exception, e:
            self.log('FATAL: could not create symlink: %s' % e)
            raise PilotException('Could not create symlink: %s' % e)

        self.log('Symlink creation successful')
        checksum, checksum_type = fspec.get_checksum()
        return {'checksum_type': checksum_type,
                'checksum': checksum,
                'filesize': fspec.filesize}

    def stageOut(self, source, destination, fspec):
        """
        Copy the output file from the pilot working directory to the destination
        directory.

        :param source:      local file location
        :param destination: remote location to copy file
        :param fspec:       dictionary containing destination replicas, scope, lfn
        :return:            destination file details (checksumtype, checksum, size)
        """

        src = os.path.realpath(fspec.lfn)
        dest = os.path.join(self.init_dir, fspec.lfn)
        self.log('Moving %s to %s' % (src, dest))

        # copy the output
        try:
            shutil.move(src, dst)
        except Exception, e:
            self.log('FATAL: could not move outputfile: %s' % e)
            raise PilotException('Could not move outputfile: %s' % e)

        self.log('Move successful')

        checksum, checksum_type = fspec.get_checksum()
        return {'checksum_type': checksum_type,
                'checksum': checksum,
                'filesize': fspec.filesize}
<commit_msg>fix stormSiteMover (move rucio import inside the function)<commit_after>"""
  storm SiteMover

  :author: Mario Lassnig <mario.lassnig@cern.ch>, 2016
"""

from .base import BaseSiteMover

from TimerCommand import TimerCommand
from PilotErrors import PilotException

from datetime import datetime
from xml.dom import minidom

import os, shutil

class stormSiteMover(BaseSiteMover):
    """
    SiteMover that queries HTTP etag for physcial storage location, then
    symlinks for stage in, and copies for stage out
    """

    name = 'storm'
    schemes = ['file']
    version = '20161125.005'

    require_replicas = False       ## quick hack to avoid query Rucio to resolve input replicas

    def __init__(self, *args, **kwargs):
        super(stormSiteMover, self).__init__(*args, **kwargs)
        self.log('storm sitemover version: %s' % self.version)

    def resolve_replica(self, fspec, protocol, ddm=None):
        """
        Overridden method -- unused
        """

        return {'ddmendpoint': fspec.replicas[0][0] if fspec.replicas else None,
                'surl': None,
                'pfn': fspec.lfn}

    def stageIn(self, source, destination, fspec):
        """
        Query HTTP for etag, then symlink to the pilot working directory.

        :param source:      original file location
        :param destination: where to create the link
        :param fspec:       dictionary containing destination replicas, scope, lfn
        :return:            destination file details (checksumtype, checksum, size)
        """

        self.log('source: %s' % str(source))
        self.log('destination: %s' % str(destination))
        self.log('fspec: %s' % str(fspec))
        self.log('fspec.scope: %s' % str(fspec.scope))
        self.log('fspec.lfn: %s' % str(fspec.lfn))
        self.log('fspec.ddmendpoint: %s' % str(fspec.ddmendpoint))

        # figure out the HTTP SURL from Rucio

        from rucio.client import ReplicaClient

        rc = ReplicaClient()
        http_surl_reps = [r for r in rc.list_replicas(dids=[{'scope': fspec.scope,
                                                             'name': fspec.lfn}],
                                                      schemes=['https'],
                                                      rse_expression=fspec.ddmendpoint)]
        self.log('http_surl_reps: %s' % http_surl_reps)

        http_surl = http_surl_reps[0]['rses'][fspec.ddmendpoint][0].rsplit('_-')[0]
        self.log('http_surl: %s' % http_surl)

        # retrieve the TURL from the webdav etag
        cmd = 'davix-http --capath /cvmfs/atlas.cern.ch/repo/ATLASLocalRootBase/etc/grid-security-emi/certificates --cert $X509_USER_PROXY -X PROPFIND %s' % http_surl
        self.log('ETAG retrieval: %s' % cmd)
        try:
            timer = TimerCommand(cmd)
            rcode, output = timer.run(timeout=10)
        except Exception, e:
            self.log('FATAL: could not retrieve STORM WebDAV ETag: %s' % e)
            raise PilotException('Could not retrieve STORM WebDAV ETag: %s' % e)
        p_output = minidom.parseString(output)

        # we need to strip off the quotation marks and the <timestamp> from the etag
        # but since we can have multiple underscores, we have to rely on the uniqueness
        # of the full LFN to make the split
        target = p_output.getElementsByTagName('d:getetag')[0].childNodes[0].nodeValue.replace('"', '')
        self.log('Symlink before: %s' % target)
        target = target.split(fspec.lfn)[0]+fspec.lfn
        self.log('Symlink after : %s' % target)

        # make the symlink
        try:
            self.log('Making symlink from %s to %s' % (target, destination))
            os.symlink(target, destination)
        except Exception, e:
            self.log('FATAL: could not create symlink: %s' % e)
            raise PilotException('Could not create symlink: %s' % e)

        self.log('Symlink creation successful')
        checksum, checksum_type = fspec.get_checksum()
        return {'checksum_type': checksum_type,
                'checksum': checksum,
                'filesize': fspec.filesize}

    def stageOut(self, source, destination, fspec):
        """
        Copy the output file from the pilot working directory to the destination
        directory.

        :param source:      local file location
        :param destination: remote location to copy file
        :param fspec:       dictionary containing destination replicas, scope, lfn
        :return:            destination file details (checksumtype, checksum, size)
        """

        src = os.path.realpath(fspec.lfn)
        dest = os.path.join(self.init_dir, fspec.lfn)
        self.log('Moving %s to %s' % (src, dest))

        # copy the output
        try:
            shutil.move(src, dst)
        except Exception, e:
            self.log('FATAL: could not move outputfile: %s' % e)
            raise PilotException('Could not move outputfile: %s' % e)

        self.log('Move successful')

        checksum, checksum_type = fspec.get_checksum()
        return {'checksum_type': checksum_type,
                'checksum': checksum,
                'filesize': fspec.filesize}
<|endoftext|>


<commit_before>import re
import warnings

import BeautifulSoup
import imdb

from django.core.management import setup_environ
from sandoval import settings
# Setting up the environment is necessary to import the models.
setup_environ(settings)
from sandoval.movie import models

class Movie(object):
    def __init__(self, title):
        # Suppress imdb warnings.
        warnings.filterwarnings('ignore', 'unable to use "lxml": No module named lxml')
        warnings.filterwarnings('ignore', 'falling back to "beautifulsoup".')
        try:
            self.find(title)
            (directors, cast, movie_model) = self.create()
            for director in directors:
                self.add_director(director, movie_model)
            for person in cast:
                self.add_cast(person, movie_model)
        except Exception, e:
            print '\t%s' % e.message
        finally:
            print 'Done.\n'

    def find(self, title):
        print 'Searching for movie "%s".' % title
        ia = imdb.IMDb()
        search_result = ia.search_movie(title, 1)
        if len(search_result) == 0:
            raise Exception, 'Movie not found!'
        self.imdb_data = search_result[0]
        self.exists(self.imdb_data.getID())
        ia.update(self.imdb_data)
        print '\tFound "%s".' % self.imdb_data['title']

    def exists(self, imdb_id):
        result = models.Movie.objects.filter(imdb_id__exact=imdb_id)
        if len(result) > 0:
            raise Exception, 'Movie already exists in database!'

    def create(self):
        movie = models.Movie()
        movie.title = self.imdb_data['title']
        plot = self.imdb_data.get('plot', ['::'])[0].split('::')
        movie.plot = plot[0]
        movie.plot_author = plot[1]
        movie.year = self.imdb_data['year']
        movie.runtime = int(re.search(
            '\d+', self.imdb_data.get('runtimes', ['0'])[0]).group(0))
        movie.rating = self.imdb_data.get('rating', 0.0)
        movie.genres = ', '.join(self.imdb_data.get('genres', []))
        movie.imdb_id = self.imdb_data.movieID
        movie.save()
        print '\tSaved movie.'
        return (self.imdb_data.get('director', []),
                self.imdb_data.get('cast', [])[:9],
                movie)

    def add_director(self, person, movie_model):
        Person(person, movie_model, True)

    def add_cast(self, person, movie_model):
        Person(person, movie_model)

class Person(object):
    def __init__(self, person, movie_model, is_director=False):
        self.model = self.exists(person.getID())
        if self.model == False:
            self.model = self.create(person.getID())
        self.model.is_director = is_director
        self.build_relation(movie_model)
        self.model.save()
        print '\tSaved person "%s".' % person['name']

    def exists(self, imdb_id):
        result = models.Person.objects.filter(imdb_id__exact=imdb_id)
        if len(result) > 0:
            return result[0]
        return False

    def create(self, imdb_id):
        ia = imdb.IMDb()
        person = ia.get_person(imdb_id)
        model = models.Person()
        print '\tImporting person "%s".' % person['name']
        name = person['canonical name'].split(', ')
        if len(name) == 2:
            model.forename = name[1]
        model.surname = name[0]
        if person.has_key('birth date'):
            model.birthdate = person['birth date']
        if person.has_key('birth notes'):
            model.birthplace = person['birth notes']
        if person.has_key('mini biography'):
            model.biography = person['mini biography'][0].split('::')[0]
        model.imdb_id = person.getID()
        return model

    def build_relation(self, movie_model):
        pass

    def save(self):
        return self.model.save()
<commit_msg>Added slug fields to insertion models.<commit_after>import re
import warnings

import BeautifulSoup
import imdb

from django.core.management import setup_environ
from django.template import defaultfilters
from sandoval import settings
# Setting up the environment is necessary to import the models.
setup_environ(settings)
from sandoval.movie import models

class Movie(object):
    def __init__(self, title):
        # Suppress imdb warnings.
        warnings.filterwarnings('ignore', 'unable to use "lxml": No module named lxml')
        warnings.filterwarnings('ignore', 'falling back to "beautifulsoup".')
        try:
            self.find(title)
            (directors, cast, movie_model) = self.create()
            for director in directors:
                self.add_director(director, movie_model)
            for person in cast:
                self.add_cast(person, movie_model)
        except Exception, e:
            print '\t%s' % e.message
        finally:
            print 'Done.\n'

    def find(self, title):
        print 'Searching for movie "%s".' % title
        ia = imdb.IMDb()
        search_result = ia.search_movie(title, 1)
        if len(search_result) == 0:
            raise Exception, 'Movie not found!'
        self.imdb_data = search_result[0]
        self.exists(self.imdb_data.getID())
        ia.update(self.imdb_data)
        print '\tFound "%s".' % self.imdb_data['title']

    def exists(self, imdb_id):
        result = models.Movie.objects.filter(imdb_id__exact=imdb_id)
        if len(result) > 0:
            raise Exception, 'Movie already exists in database!'

    def create(self):
        movie = models.Movie()
        movie.title = self.imdb_data['title']
        movie.slug = defaultfilters.slugify(self.imdb_data['title'] + ' '
            + self.imdb_data['year'].__str__())
        plot = self.imdb_data.get('plot', ['::'])[0].split('::')
        movie.plot = plot[0]
        movie.plot_author = plot[1]
        movie.year = self.imdb_data['year']
        movie.runtime = int(re.search(
            '\d+', self.imdb_data.get('runtimes', ['0'])[0]).group(0))
        movie.rating = self.imdb_data.get('rating', 0.0)
        movie.genres = ', '.join(self.imdb_data.get('genres', []))
        movie.imdb_id = self.imdb_data.movieID
        movie.save()
        print '\tSaved movie.'
        return (self.imdb_data.get('director', []),
                self.imdb_data.get('cast', [])[:9],
                movie)

    def add_director(self, person, movie_model):
        Person(person, movie_model, True)

    def add_cast(self, person, movie_model):
        Person(person, movie_model)

class Person(object):
    def __init__(self, person, movie_model, is_director=False):
        self.model = self.exists(person.getID())
        if self.model == False:
            self.model = self.create(person.getID())
        self.model.is_director = is_director
        self.build_relation(movie_model)
        self.model.save()
        print '\tSaved person "%s".' % person['name']

    def exists(self, imdb_id):
        result = models.Person.objects.filter(imdb_id__exact=imdb_id)
        if len(result) > 0:
            return result[0]
        return False

    def create(self, imdb_id):
        ia = imdb.IMDb()
        person = ia.get_person(imdb_id)
        model = models.Person()
        print '\tImporting person "%s".' % person['name']
        name = person['canonical name'].split(', ')
        if len(name) == 2:
            model.forename = name[1]
        model.surname = name[0]
        try:
            # Try to set the slug.
            model.slug = defaultfilters.slugify(person['name'])
        except IntergrityError:
            # Setting the slug failed because the value was not unique.
            model.slug = defaultfilters.slugify(person['name'] + ' ' +
                person.getID())
        if person.has_key('birth date'):
            model.birthdate = person['birth date']
        if person.has_key('birth notes'):
            model.birthplace = person['birth notes']
        if person.has_key('mini biography'):
            model.biography = person['mini biography'][0].split('::')[0]
        model.imdb_id = person.getID()
        return model

    def build_relation(self, movie_model):
        pass

    def save(self):
        return self.model.save()
<|endoftext|>


<commit_before>"""Contains the ShotGroup base class."""

from collections import deque

from mpf.core.mode_device import ModeDevice
from mpf.core.system_wide_device import SystemWideDevice
from mpf.core.utility_functions import Util

class ShotGroup(ModeDevice, SystemWideDevice):

    """Represents a group of shots in a pinball machine by grouping together multiple `Shot` class devices.

    This is used so you get get
    "group-level" functionality, like shot rotation, shot group completion,
    etc. This would be used for a group of rollover lanes, a bank of standups,
    etc.
    """

    config_section = 'shot_groups'
    collection = 'shot_groups'
    class_label = 'shot_group'

    def __init__(self, machine, name):
        """Initialise shot group."""
        super().__init__(machine, name)

        self.rotation_enabled = True
        self._enabled = False

        # todo remove this hack
        self._created_system_wide = False

    @property
    def enabled(self):
        """Return true if enabled."""
        return self._enabled

    @classmethod
    def prepare_config(cls, config, is_mode_config):
        """Add default events if not in mode."""
        if not is_mode_config:
            # If this device is setup in a machine-wide config, make sure it has
            # a default enable event.
            if 'enable_events' not in config:
                config['enable_events'] = 'ball_starting'
            if 'disable_events' not in config:
                config['disable_events'] = 'ball_ended'
            if 'reset_events' not in config:
                config['reset_events'] = 'ball_ended'
        return config

    def device_added_system_wide(self):
        """Called when a device is added system wide."""
        super().device_added_system_wide()

        self._created_system_wide = True

    def device_added_to_mode(self, mode, player):
        """Add device in mode."""
        super().device_added_to_mode(mode, player)

        # If there are no enable_events configured, then we enable this shot
        # group when its created on mode start

        if ((not mode.config['shot_groups'][self.name]['enable_events']) or
                'mode_{}_started'.format(mode.name) in
                mode.config['shot_groups'][self.name]['enable_events']):
            self.enable(mode)
        else:
            # manually call disable here so it disables the member shots
            self.disable(mode)

    def _enable_related_device_debugging(self):
        self.log.debug(
            "Enabling debugging for this shot groups's member shots")

        for shot in self.config['shots']:
            shot.enable_debugging()

    def _disable_related_device_debugging(self):
        for shot in self.config['shots']:
            shot.disable_debugging()

    def hit(self, mode, profile, state, **kwargs):
        """One of the member shots in this shot group was hit.

        Args:
            profile: String name of the active profile of the shot that
                was hit.
            mode: unused
            kwargs: unused
        """
        del mode
        del kwargs

        if not self._enabled:
            return

        self.debug_log('Hit! Active profile: %s, Current state: %s',
                       profile, profile)

        self.machine.events.post(self.name + '_hit',
                                 profile=profile, state=state)
        '''event: (shot_group)_hit
        desc: A member shot in the shot group called (shot_group) was just hit.

        Note that there are three events posted when a member shot is hit, each
        with variants of the shot name, profile, and current state,
        allowing you to key in on the specific granularity you need.

        Also remember that shots can have more than one active profile at a
        time (typically each associated with a mode), so a single hit to this
        shot might result in this event being posted multiple times with
        different (profile) values.

        args:
        profile: The name of the profile that was active when hit.
        state: The name of the state the profile was in when it was hit
        '''

        self.machine.events.post(self.name + '_' + profile + '_hit',
                                 profile=profile, state=state)
        '''event: (shot_group)_(profile)_hit
        desc: A member shot in the shot group called (shot_group) was just hit
        with the profile called (profile) applied.

        Note that there are three events posted when a member shot is hit, each
        with variants of the shot name, profile, and current state,
        allowing you to key in on the specific granularity you need.

        Also remember that shots can have more than one active profile at a
        time (typically each associated with a mode), so a single hit to this
        shot might result in this event being posted multiple times with
        different (profile) values.

        args:
        profile: The name of the profile that was active when hit.
        state: The name of the state the profile was in when it was hit'''
        self.machine.events.post(self.name + '_' + profile + '_' + state +
                                 '_hit', profile=profile, state=state)
        '''event: (shot_group)_(profile)_(state)_hit
        desc: A member shot in the shot group called (shot_group) was just hit
        with the profile called (profile) applied in the current state called
        (state).

        Note that there are three events posted when a member shot is hit, each
        with variants of the shot name, profile, and current state,
        allowing you to key in on the specific granularity you need.

        Also remember that shots can have more than one active profile at a
        time (typically each associated with a mode), so a single hit to this
        shot might result in this event being posted multiple times with
        different (profile) values.

        args:
        profile: The name of the profile that was active when hit.
        state: The name of the state the profile was in when it was hit'''
    def enable(self, mode=None, profile=None, **kwargs):
        """Enable this shot group.

        Also enables all the shots in this group.
        """
        del kwargs

        if mode:
            self._enable_from_mode(mode, profile)
        else:
            self._enable_from_system_wide(profile)

        self._enabled = True

        self.rotation_enabled = not self.config['enable_rotation_events']

        self.debug_log('Enabling from mode: %s', mode)

    def _enable_from_mode(self, mode, profile=None):
        # If we weren't passed a profile, use the one from the mode config

        if not profile and mode.config['shot_groups'][self.name]['profile']:
            profile = mode.config['shot_groups'][self.name]['profile']

        for shot in self.config['shots']:
            if profile:
                # this is a passed profile or the profile in the shot group
                # in the mode config
                shot.update_profile(profile=profile, mode=mode, enable=True)

            else:
                shot.update_profile(mode=mode, enable=True)

            shot.register_group(self)

    def _enable_from_system_wide(self, profile=None):
        for shot in self.config['shots']:
            if 'profile' in self.config:
                shot.update_profile(profile=self.config['profile'],
                                    mode=None)
            shot.enable(profile=profile)
            shot.register_group(self)

    def disable(self, mode=None, **kwargs):
        """Disable this shot group.

        Also disables all the shots in this group.
        """
        del kwargs
        self._enabled = False
        self.debug_log('Disabling from mode: %s', mode)

        for shot in self.config['shots']:
            shot.disable(mode)
            shot.deregister_group(self)

    def enable_rotation(self, **kwargs):
        """Enable shot rotation.

        If disabled, rotation events do not actually rotate the shots.
        """
        del kwargs
        self.debug_log('Enabling rotation')
        self.rotation_enabled = True

    def disable_rotation(self, **kwargs):
        """Disable shot rotation.

        If disabled, rotation events do not actually rotate the shots.
        """
        del kwargs
        self.debug_log('Disabling rotation')
        self.rotation_enabled = False

    def reset(self, mode=None, **kwargs):
        """Reset each of the shots in this group back to the initial state in whatever shot profile they have applied.

        This is the same as calling each shot's reset() method one-by-one.
        """
        del kwargs
        self.debug_log('Resetting')
        for shot in self.config['shots']:
            shot.reset(mode)

    def remove_active_profile(self, mode, **kwargs):
        """Remove the current active profile from every shot in the group."""
        del kwargs
        self.debug_log('Removing active profile')
        for shot in self.config['shots']:
            shot.remove_active_profile(mode)

    def advance(self, steps=1, mode=None, force=False, **kwargs):
        """Advance the current active profile from every shot in the group one step forward."""
        del kwargs

        if not (self._enabled or force):
            return

        self.debug_log('Advancing')
        for shot in self.config['shots']:
            shot.advance(steps=steps, mode=mode, force=force)

    def rotate(self, direction=None, states=None,
               exclude_states=None, mode=None, **kwargs):
        """Rotate (or "shift") the state of all the shots in this group.

        This is used for things like lane change, where hitting the flipper
        button shifts all the states of the shots in the group to the left or
        right.

        This method actually transfers the current state of each shot profile
        to the left or the right, and the shot on the end rolls over to the
        taret on the other end.

        Args:
            direction: String that specifies whether the rotation direction is
                to the left or right. Values are 'right' or 'left'. Default of
                None will cause the shot group to rotate in the direction as
                specified by the rotation_pattern.
            states: A string of a state or a list of strings that represent the
                targets that will be selected to rotate. If None (default), then
                all targets will be included.
            exclude_states: A string of a state or a list of strings that
                controls whether any targets will *not* be rotated. (Any
                targets with an active profile in one of these states will not
                be included in the rotation. Default is None which means all
                targets will be rotated)
            kwargs: unused

        Note that this shot group must, and rotation_events for this
        shot group, must both be enabled for the rotation events to work.
        """
        del kwargs

        if not (self._enabled and self.rotation_enabled):
            self.debug_log("Received rotation request. "
                           "Rotation Enabled: %s. Will NOT rotate",
                           self.rotation_enabled)

            return

        # if we don't have states or exclude_states, we'll see if the first shot
        # in the group has them and use those. Since all the shots should have
        # the same profile applied, it's ok to just pick from the first one.

        if states:
            states = Util.string_to_lowercase_list(states)
        else:
            states = self.config['shots'][0].get_profile_by_key(
                'mode', mode)['settings']['state_names_to_rotate']

        if exclude_states:
            exclude_states = Util.string_to_lowercase_list(exclude_states)
        else:
            exclude_states = (
                self.config['shots'][0].get_profile_by_key(
                    'mode', mode)['settings']['state_names_to_not_rotate'])

        shot_list = list()

        # build of a list of shots we're actually going to rotate
        for shot in self.config['shots']:

            if ((not states or shot.get_profile_by_key(
                    'mode', mode)['current_state_name'] in states) and
                    shot.get_profile_by_key(
                    'mode', mode)['current_state_name'] not in exclude_states):
                shot_list.append(shot)

        # shot_state_list is deque of tuples (state num, show step num)
        shot_state_list = deque()

        for shot in shot_list:
            try:
                current_show_step = shot.get_profile_by_key('mode', mode)['running_show'].next_step_index
            except AttributeError:
                current_show_step = None

            shot_state_list.append(
                (shot.player[shot.get_profile_by_key('mode', mode)['settings']['player_variable']], current_show_step))

        if self.debug:
            self.log.debug('Rotating. Mode: %s, Direction: %s, Include states:'
                           ' %s, Exclude states: %s, Shots to be rotated: %s',
                           mode, direction, states,
                           exclude_states, [x.name for x in shot_list])

            for shot in shot_list:
                shot.log.debug("This shot is part of a rotation event. Current"
                               " state: %s", shot.get_profile_by_key('mode', mode)['current_state_name'])

        # figure out which direction we're going to rotate
        if not direction:
            direction = shot_list[0].get_profile_by_key(
                'mode', mode)['settings']['rotation_pattern'][0]
            shot_list[0].get_profile_by_key(
                'mode', mode)['settings']['rotation_pattern'].rotate(-1)

            self.debug_log("Since no direction was specified, pulling from"
                           " rotation pattern: '%s'", direction)

        # rotate that list
        if direction == 'right':
            shot_state_list.rotate(1)
        else:
            shot_state_list.rotate(-1)

        # step through all our shots and update their states
        for i, shot in enumerate(shot_list):
            shot.jump(mode=mode, state=shot_state_list[i][0],
                      show_step=shot_state_list[i][1], force=True)

    def rotate_right(self, mode=None, **kwargs):
        """Rotate the state of the shots to the right.

        This method is the same as calling rotate('right')

        Args:
            kwargs: unused

        """
        del kwargs
        self.rotate(direction='right', mode=mode)

    def rotate_left(self, mode=None, **kwargs):
        """Rotate the state of the shots to the left.

        This method is the same as calling rotate('left')

        Args:
            kwargs: unused

        """
        del kwargs
        self.rotate(direction='left', mode=mode)

    def check_for_complete(self, mode):
        """Check all the shots in this shot group.

        If they are all in the same state, then a complete event is posted.
        """
        # TODO should be made to work for lower priority things too?
        if not self._enabled:
            return

        shot_states = set()

        self.debug_log("Checking for complete. mode: %s", mode)

        for shot in self.config['shots']:

            mode_state = shot.get_profile_by_key('mode', mode)

            if mode_state:
                shot_states.add((mode_state['profile'],
                                 mode_state['current_state_name']))

            else:
                self.debug_log("Shot %s is not used in this mode. Aborting"
                               " check for complete", shot)
                return

        if len(shot_states) == 1:
            profile, state = shot_states.pop()

            self.debug_log(
                "Shot group is complete with profile :%s, state:"
                "%s", profile, state)

            self.machine.events.post(self.name + '_complete')
            '''event: (shot_group)_complete
            desc: All the member shots in the shot group called (shot_group)
            are in the same state.
            '''

            self.machine.events.post(self.name + '_' + profile + '_complete')
            '''event: (shot_group)_(profile)_complete
            desc: All the member shots in the shot group called (shot_group)
            with the profile called (profile) are in the same state.
            '''

            self.machine.events.post(self.name + '_' + profile + '_' + state +
                                     '_complete')
            '''event: (shot_group)_(profile)_(state)_complete
            desc: All the member shots in the shot group called (shot_group)
            with the profile called (profile) are in the same state with the
            name (state).
            '''

    def add_control_events_in_mode(self, mode):
        """Add control events in mode."""
        # called if any control_events for this shot_group exist in the mode
        # config, regardless of whether or not the shot_group device was
        # initially created in this mode
        for shot in self.config['shots']:
            if not shot.get_profile_by_key('mode', mode):

                self.debug_log('Control events found in %s '
                               'config. Adding profile list entries to '
                               'member shots', mode)

                enable = not self.config['enable_events']

                if self.config['profile']:
                    profile = self.config['profile']
                else:
                    profile = shot.config['profile']

                shot.update_profile(profile=profile, enable=enable, mode=mode)

    def device_removed_from_mode(self, mode):
        """Disable device when mode stops."""
        del mode
        if self._created_system_wide:
            return
        self.debug_log("Removing this shot group")
        self._enabled = False
<commit_msg>better error<commit_after>"""Contains the ShotGroup base class."""

from collections import deque

from mpf.core.mode_device import ModeDevice
from mpf.core.system_wide_device import SystemWideDevice
from mpf.core.utility_functions import Util

class ShotGroup(ModeDevice, SystemWideDevice):

    """Represents a group of shots in a pinball machine by grouping together multiple `Shot` class devices.

    This is used so you get get
    "group-level" functionality, like shot rotation, shot group completion,
    etc. This would be used for a group of rollover lanes, a bank of standups,
    etc.
    """

    config_section = 'shot_groups'
    collection = 'shot_groups'
    class_label = 'shot_group'

    def __init__(self, machine, name):
        """Initialise shot group."""
        super().__init__(machine, name)

        self.rotation_enabled = True
        self._enabled = False

        # todo remove this hack
        self._created_system_wide = False

    @property
    def enabled(self):
        """Return true if enabled."""
        return self._enabled

    @classmethod
    def prepare_config(cls, config, is_mode_config):
        """Add default events if not in mode."""
        if not is_mode_config:
            # If this device is setup in a machine-wide config, make sure it has
            # a default enable event.
            if 'enable_events' not in config:
                config['enable_events'] = 'ball_starting'
            if 'disable_events' not in config:
                config['disable_events'] = 'ball_ended'
            if 'reset_events' not in config:
                config['reset_events'] = 'ball_ended'
        return config

    def device_added_system_wide(self):
        """Called when a device is added system wide."""
        super().device_added_system_wide()

        self._created_system_wide = True

    def device_added_to_mode(self, mode, player):
        """Add device in mode."""
        super().device_added_to_mode(mode, player)

        # If there are no enable_events configured, then we enable this shot
        # group when its created on mode start

        if ((not mode.config['shot_groups'][self.name]['enable_events']) or
                'mode_{}_started'.format(mode.name) in
                mode.config['shot_groups'][self.name]['enable_events']):
            self.enable(mode)
        else:
            # manually call disable here so it disables the member shots
            self.disable(mode)

    def _enable_related_device_debugging(self):
        self.log.debug(
            "Enabling debugging for this shot groups's member shots")

        for shot in self.config['shots']:
            shot.enable_debugging()

    def _disable_related_device_debugging(self):
        for shot in self.config['shots']:
            shot.disable_debugging()

    def hit(self, mode, profile, state, **kwargs):
        """One of the member shots in this shot group was hit.

        Args:
            profile: String name of the active profile of the shot that
                was hit.
            mode: unused
            kwargs: unused
        """
        del mode
        del kwargs

        if not profile:
            raise AssertionError("Called hit without profile")

        if not self._enabled:
            return

        self.debug_log('Hit! Active profile: %s, Current state: %s',
                       profile, profile)

        self.machine.events.post(self.name + '_hit',
                                 profile=profile, state=state)
        '''event: (shot_group)_hit
        desc: A member shot in the shot group called (shot_group) was just hit.

        Note that there are three events posted when a member shot is hit, each
        with variants of the shot name, profile, and current state,
        allowing you to key in on the specific granularity you need.

        Also remember that shots can have more than one active profile at a
        time (typically each associated with a mode), so a single hit to this
        shot might result in this event being posted multiple times with
        different (profile) values.

        args:
        profile: The name of the profile that was active when hit.
        state: The name of the state the profile was in when it was hit
        '''

        self.machine.events.post(self.name + '_' + profile + '_hit',
                                 profile=profile, state=state)
        '''event: (shot_group)_(profile)_hit
        desc: A member shot in the shot group called (shot_group) was just hit
        with the profile called (profile) applied.

        Note that there are three events posted when a member shot is hit, each
        with variants of the shot name, profile, and current state,
        allowing you to key in on the specific granularity you need.

        Also remember that shots can have more than one active profile at a
        time (typically each associated with a mode), so a single hit to this
        shot might result in this event being posted multiple times with
        different (profile) values.

        args:
        profile: The name of the profile that was active when hit.
        state: The name of the state the profile was in when it was hit'''
        self.machine.events.post(self.name + '_' + profile + '_' + state +
                                 '_hit', profile=profile, state=state)
        '''event: (shot_group)_(profile)_(state)_hit
        desc: A member shot in the shot group called (shot_group) was just hit
        with the profile called (profile) applied in the current state called
        (state).

        Note that there are three events posted when a member shot is hit, each
        with variants of the shot name, profile, and current state,
        allowing you to key in on the specific granularity you need.

        Also remember that shots can have more than one active profile at a
        time (typically each associated with a mode), so a single hit to this
        shot might result in this event being posted multiple times with
        different (profile) values.

        args:
        profile: The name of the profile that was active when hit.
        state: The name of the state the profile was in when it was hit'''
    def enable(self, mode=None, profile=None, **kwargs):
        """Enable this shot group.

        Also enables all the shots in this group.
        """
        del kwargs

        if mode:
            self._enable_from_mode(mode, profile)
        else:
            self._enable_from_system_wide(profile)

        self._enabled = True

        self.rotation_enabled = not self.config['enable_rotation_events']

        self.debug_log('Enabling from mode: %s', mode)

    def _enable_from_mode(self, mode, profile=None):
        # If we weren't passed a profile, use the one from the mode config

        if not profile and mode.config['shot_groups'][self.name]['profile']:
            profile = mode.config['shot_groups'][self.name]['profile']

        for shot in self.config['shots']:
            if profile:
                # this is a passed profile or the profile in the shot group
                # in the mode config
                shot.update_profile(profile=profile, mode=mode, enable=True)

            else:
                shot.update_profile(mode=mode, enable=True)

            shot.register_group(self)

    def _enable_from_system_wide(self, profile=None):
        for shot in self.config['shots']:
            if 'profile' in self.config:
                shot.update_profile(profile=self.config['profile'],
                                    mode=None)
            shot.enable(profile=profile)
            shot.register_group(self)

    def disable(self, mode=None, **kwargs):
        """Disable this shot group.

        Also disables all the shots in this group.
        """
        del kwargs
        self._enabled = False
        self.debug_log('Disabling from mode: %s', mode)

        for shot in self.config['shots']:
            shot.disable(mode)
            shot.deregister_group(self)

    def enable_rotation(self, **kwargs):
        """Enable shot rotation.

        If disabled, rotation events do not actually rotate the shots.
        """
        del kwargs
        self.debug_log('Enabling rotation')
        self.rotation_enabled = True

    def disable_rotation(self, **kwargs):
        """Disable shot rotation.

        If disabled, rotation events do not actually rotate the shots.
        """
        del kwargs
        self.debug_log('Disabling rotation')
        self.rotation_enabled = False

    def reset(self, mode=None, **kwargs):
        """Reset each of the shots in this group back to the initial state in whatever shot profile they have applied.

        This is the same as calling each shot's reset() method one-by-one.
        """
        del kwargs
        self.debug_log('Resetting')
        for shot in self.config['shots']:
            shot.reset(mode)

    def remove_active_profile(self, mode, **kwargs):
        """Remove the current active profile from every shot in the group."""
        del kwargs
        self.debug_log('Removing active profile')
        for shot in self.config['shots']:
            shot.remove_active_profile(mode)

    def advance(self, steps=1, mode=None, force=False, **kwargs):
        """Advance the current active profile from every shot in the group one step forward."""
        del kwargs

        if not (self._enabled or force):
            return

        self.debug_log('Advancing')
        for shot in self.config['shots']:
            shot.advance(steps=steps, mode=mode, force=force)

    def rotate(self, direction=None, states=None,
               exclude_states=None, mode=None, **kwargs):
        """Rotate (or "shift") the state of all the shots in this group.

        This is used for things like lane change, where hitting the flipper
        button shifts all the states of the shots in the group to the left or
        right.

        This method actually transfers the current state of each shot profile
        to the left or the right, and the shot on the end rolls over to the
        taret on the other end.

        Args:
            direction: String that specifies whether the rotation direction is
                to the left or right. Values are 'right' or 'left'. Default of
                None will cause the shot group to rotate in the direction as
                specified by the rotation_pattern.
            states: A string of a state or a list of strings that represent the
                targets that will be selected to rotate. If None (default), then
                all targets will be included.
            exclude_states: A string of a state or a list of strings that
                controls whether any targets will *not* be rotated. (Any
                targets with an active profile in one of these states will not
                be included in the rotation. Default is None which means all
                targets will be rotated)
            kwargs: unused

        Note that this shot group must, and rotation_events for this
        shot group, must both be enabled for the rotation events to work.
        """
        del kwargs

        if not (self._enabled and self.rotation_enabled):
            self.debug_log("Received rotation request. "
                           "Rotation Enabled: %s. Will NOT rotate",
                           self.rotation_enabled)

            return

        # if we don't have states or exclude_states, we'll see if the first shot
        # in the group has them and use those. Since all the shots should have
        # the same profile applied, it's ok to just pick from the first one.

        if states:
            states = Util.string_to_lowercase_list(states)
        else:
            states = self.config['shots'][0].get_profile_by_key(
                'mode', mode)['settings']['state_names_to_rotate']

        if exclude_states:
            exclude_states = Util.string_to_lowercase_list(exclude_states)
        else:
            exclude_states = (
                self.config['shots'][0].get_profile_by_key(
                    'mode', mode)['settings']['state_names_to_not_rotate'])

        shot_list = list()

        # build of a list of shots we're actually going to rotate
        for shot in self.config['shots']:

            if ((not states or shot.get_profile_by_key(
                    'mode', mode)['current_state_name'] in states) and
                    shot.get_profile_by_key(
                    'mode', mode)['current_state_name'] not in exclude_states):
                shot_list.append(shot)

        # shot_state_list is deque of tuples (state num, show step num)
        shot_state_list = deque()

        for shot in shot_list:
            try:
                current_show_step = shot.get_profile_by_key('mode', mode)['running_show'].next_step_index
            except AttributeError:
                current_show_step = None

            shot_state_list.append(
                (shot.player[shot.get_profile_by_key('mode', mode)['settings']['player_variable']], current_show_step))

        if self.debug:
            self.log.debug('Rotating. Mode: %s, Direction: %s, Include states:'
                           ' %s, Exclude states: %s, Shots to be rotated: %s',
                           mode, direction, states,
                           exclude_states, [x.name for x in shot_list])

            for shot in shot_list:
                shot.log.debug("This shot is part of a rotation event. Current"
                               " state: %s", shot.get_profile_by_key('mode', mode)['current_state_name'])

        # figure out which direction we're going to rotate
        if not direction:
            direction = shot_list[0].get_profile_by_key(
                'mode', mode)['settings']['rotation_pattern'][0]
            shot_list[0].get_profile_by_key(
                'mode', mode)['settings']['rotation_pattern'].rotate(-1)

            self.debug_log("Since no direction was specified, pulling from"
                           " rotation pattern: '%s'", direction)

        # rotate that list
        if direction == 'right':
            shot_state_list.rotate(1)
        else:
            shot_state_list.rotate(-1)

        # step through all our shots and update their states
        for i, shot in enumerate(shot_list):
            shot.jump(mode=mode, state=shot_state_list[i][0],
                      show_step=shot_state_list[i][1], force=True)

    def rotate_right(self, mode=None, **kwargs):
        """Rotate the state of the shots to the right.

        This method is the same as calling rotate('right')

        Args:
            kwargs: unused

        """
        del kwargs
        self.rotate(direction='right', mode=mode)

    def rotate_left(self, mode=None, **kwargs):
        """Rotate the state of the shots to the left.

        This method is the same as calling rotate('left')

        Args:
            kwargs: unused

        """
        del kwargs
        self.rotate(direction='left', mode=mode)

    def check_for_complete(self, mode):
        """Check all the shots in this shot group.

        If they are all in the same state, then a complete event is posted.
        """
        # TODO should be made to work for lower priority things too?
        if not self._enabled:
            return

        shot_states = set()

        self.debug_log("Checking for complete. mode: %s", mode)

        for shot in self.config['shots']:

            mode_state = shot.get_profile_by_key('mode', mode)

            if mode_state:
                shot_states.add((mode_state['profile'],
                                 mode_state['current_state_name']))

            else:
                self.debug_log("Shot %s is not used in this mode. Aborting"
                               " check for complete", shot)
                return

        if len(shot_states) == 1:
            profile, state = shot_states.pop()

            self.debug_log(
                "Shot group is complete with profile :%s, state:"
                "%s", profile, state)

            self.machine.events.post(self.name + '_complete')
            '''event: (shot_group)_complete
            desc: All the member shots in the shot group called (shot_group)
            are in the same state.
            '''

            self.machine.events.post(self.name + '_' + profile + '_complete')
            '''event: (shot_group)_(profile)_complete
            desc: All the member shots in the shot group called (shot_group)
            with the profile called (profile) are in the same state.
            '''

            self.machine.events.post(self.name + '_' + profile + '_' + state +
                                     '_complete')
            '''event: (shot_group)_(profile)_(state)_complete
            desc: All the member shots in the shot group called (shot_group)
            with the profile called (profile) are in the same state with the
            name (state).
            '''

    def add_control_events_in_mode(self, mode):
        """Add control events in mode."""
        # called if any control_events for this shot_group exist in the mode
        # config, regardless of whether or not the shot_group device was
        # initially created in this mode
        for shot in self.config['shots']:
            if not shot.get_profile_by_key('mode', mode):

                self.debug_log('Control events found in %s '
                               'config. Adding profile list entries to '
                               'member shots', mode)

                enable = not self.config['enable_events']

                if self.config['profile']:
                    profile = self.config['profile']
                else:
                    profile = shot.config['profile']

                shot.update_profile(profile=profile, enable=enable, mode=mode)

    def device_removed_from_mode(self, mode):
        """Disable device when mode stops."""
        del mode
        if self._created_system_wide:
            return
        self.debug_log("Removing this shot group")
        self._enabled = False
<|endoftext|>