<commit_before>import apscheduler, os, redis, string, subprocess
from datetime import timedelta
from multiprocessing import Pool
from plumbum import SshMachine
from rpyc.utils.factory import ssh_connect
from utilities import distance, threaded_map

process_pool_size = 60
FNULL = open(os.devnull, 'w')

all_dns = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=0))
open_resolvers = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=1))
geoip = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=2))
pl_hosts = [line.split(' ')[0:4] for line in map(string.strip,open('pl-host-list-geo').readlines())]

class PlanetLabNode(object):
    def __init__(self, host, ip, lat, lon):
        self.host = host
        self.ip = ip
        self.lat = float(lat)
        self.lon = float(lon)
        self.connected = False

    def connect(self):
        try:
            self.connectPL()
        except (AssertionError, EOFError), e:
            try:
                self.restartPL()
                self.connectPL()
            except (AssertionError, EOFError), e:
                self.connected = False

    def connectPL(self):
        rem = SshMachine(self.host, user='ucb_268_measure', keyfile='~/.ssh/id_rsa',
                         ssh_opts=["-o StrictHostKeyChecking no",
                                   "-o UserKnownHostsFile=/dev/null"])
        conn = ssh_connect(rem, 18861, config={'allow_pickle' : True})
        assert conn.root.test() == 1, "%s: RPC Failure" % self.host
        self.conn = conn
        self.connected = True

    def restartPL(self):
        subprocess.call(["ssh", "-t", "-i", "~/.ssh/id_rsa", "-o StrictHostKeyChecking no",
                         "-o UserKnownHostsFile=/dev/null", "ucb_268_measure@%s" % self.host,
                         "sudo tking-server stop; sudo tking-server start;"],
                         stdout=FNULL, stderr=subprocess.STDOUT)

    def handleConnExceptions(fn):
        def wrapped(self, *args, **kwargs):
            if not self.connected:
                self.connect()

            if self.connected:
                try:
                    return fn(self, *args, **kwargs)
                except Exception, e:
                    self.connected = False
                    try:
                        self.connect()
                        return fn(self, *args, **kwargs)
                    except Exception, e:
                        self.connected = False
                        return None
            else:
                return None
        return wrapped

    @handleConnExceptions
    def get_latency(self, target1, ip1, target2, ip2):
        return self.conn.root.get_latency(target1, ip1, target2, ip2)

    @handleConnExceptions
    def get_k(self, target, ip):
        return self.conn.get_k(target, ip)

    @handleConnExceptions
    def get_distance(self, lat, lon):
        return distance((self.lat, self.lon), (lat, lon))

def select_random_points():
    target1 = open_resolvers.randomkey()
    target2 = geoip.randomkey()
    while not geoip.exists(target1):
        target1 = open_resolvers.randomkey()

    ip1, coord1 = list(all_dns.smembers(target1))[0], eval(list(geoip.smembers(target1))[0])[1:]
    ip2, coord2 = list(all_dns.smembers(target2))[0], eval(list(geoip.smembers(target2))[0])[1:]

    return (target1, ip1, coord1), (target2, ip2, coord2)

def query_latency(target1, target2):
    name1, ip1, coord1 = target1
    name2, ip2, coord2 = target2

    # Get closest 4 PL Nodes
    distances = map(lambda node: (distance(coord1, (node.lat, node.lon)), node), pl_nodes)
    distances.sort()
    distances = distances[:4]

    results = threaded_map(lambda (dist, node): node.get_latency(name1, ip1, name2, ip2), distances, timeout=10.0)
    return results

def one_round(x):
    return query_latency(*select_random_points())

pl_nodes = map(lambda args: PlanetLabNode(*args), pl_hosts)
p = Pool(process_pool_size)
results = p.map(one_round, range(process_pool_size))

for r in results:
    end_time, start_time, ping_time, address = r
    print ping_time
#    if ping_time:
#        avg_ping = sum(filter(None,ping_time), timedelta())/len(ping_time)
#    else:
#        avg_ping = None
    print address, end_time, start_time

<commit_msg>fixed broken pipe error<commit_after>import apscheduler, os, redis, string, subprocess
from datetime import timedelta
from multiprocessing import Pool
from plumbum import SshMachine
from rpyc.utils.factory import ssh_connect
from utilities import distance, threaded_map

process_pool_size = 60
FNULL = open(os.devnull, 'w')

all_dns = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=0))
open_resolvers = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=1))
geoip = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=2))
pl_hosts = [line.split(' ')[0:4] for line in map(string.strip,open('pl-host-list-geo').readlines())]

class PlanetLabNode(object):
    def __init__(self, host, ip, lat, lon):
        self.host = host
        self.ip = ip
        self.lat = float(lat)
        self.lon = float(lon)
        self.connected = False

    def connect(self):
        try:
            self.connectPL()
        except (AssertionError, EOFError), e:
            try:
                self.restartPL()
                self.connectPL()
            except (AssertionError, EOFError), e:
                self.connected = False

    def connectPL(self):
        rem = SshMachine(self.host, user='ucb_268_measure', keyfile='~/.ssh/id_rsa',
                         ssh_opts=["-o StrictHostKeyChecking no",
                                   "-o UserKnownHostsFile=/dev/null"])
        conn = ssh_connect(rem, 18861, config={'allow_pickle' : True})
        assert conn.root.test() == 1, "%s: RPC Failure" % self.host
        self.conn = conn
        self.connected = True

    def restartPL(self):
        print 'Restart'
        subprocess.call(["ssh", "-t", "-i", "~/.ssh/id_rsa", "-o StrictHostKeyChecking no",
                         "-o UserKnownHostsFile=/dev/null", "ucb_268_measure@%s" % self.host,
                         "sudo tking-server stop; sudo tking-server start;"])

    def handleConnExceptions(fn):
        def wrapped(self, *args, **kwargs):
            if not self.connected:
                self.connect()

            if self.connected:
                try:
                    return fn(self, *args, **kwargs)
                except Exception, e:
                    self.connected = False
                    try:
                        self.connect()
                        return fn(self, *args, **kwargs)
                    except Exception, e:
                        self.connected = False
                        return None
            else:
                return None
        return wrapped

    @handleConnExceptions
    def get_latency(self, target1, ip1, target2, ip2):
        return self.conn.root.get_latency(target1, ip1, target2, ip2)

    @handleConnExceptions
    def get_k(self, target, ip):
        return self.conn.get_k(target, ip)

    @handleConnExceptions
    def get_distance(self, lat, lon):
        return distance((self.lat, self.lon), (lat, lon))

def select_random_points():
    target1 = open_resolvers.randomkey()
    target2 = geoip.randomkey()
    while not geoip.exists(target1):
        target1 = open_resolvers.randomkey()

    ip1, coord1 = list(all_dns.smembers(target1))[0], eval(list(geoip.smembers(target1))[0])[1:]
    ip2, coord2 = list(all_dns.smembers(target2))[0], eval(list(geoip.smembers(target2))[0])[1:]

    return (target1, ip1, coord1), (target2, ip2, coord2)

def query_latency(target1, target2):
    name1, ip1, coord1 = target1
    name2, ip2, coord2 = target2

    # Get closest 4 PL Nodes
    distances = map(lambda node: (distance(coord1, (node.lat, node.lon)), node), pl_nodes)
    distances.sort()
    distances = distances[:4]

    results = threaded_map(lambda (dist, node): node.get_latency(name1, ip1, name2, ip2), distances, timeout=10.0)
    return results

def one_round(x):
    return query_latency(*select_random_points())

pl_nodes = map(lambda args: PlanetLabNode(*args), pl_hosts)
p = Pool(process_pool_size)
results = p.map(one_round, range(process_pool_size))

for r in results:
    end_time, start_time, ping_time, address = r
    print ping_time
#    if ping_time:
#        avg_ping = sum(filter(None,ping_time), timedelta())/len(ping_time)
#    else:
#        avg_ping = None
    print address, end_time, start_time

<|endoftext|>


<commit_before>import logging, os, redis, rpyc, string, subprocess
from apscheduler.scheduler import Scheduler
from datetime import datetime, timedelta
from multiprocessing import Process
from plumbum import SshMachine
from rpyc.utils.factory import ssh_connect
from utilities import distance

process_pool_size = 50

all_dns = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=0))
open_resolvers = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=1))
geoip = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=2))
pl_hosts = [line.split(' ')[0:4] for line in map(string.strip,open('pl-host-list-geo').readlines())]

from sqlalchemy import create_engine
from sqlalchemy import Column, Integer, PickleType
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base

engine = create_engine('sqlite:///latency.db', echo=False)
Base = declarative_base(bind=engine)

class DataPoint(object):
    def __init__(self, start, end, pings, address, test_point):
        self.start = start
        self.end = end
        self.pings = pings
        self.address = address
        self.test_point = test_point

class Query(Base):
    __tablename__ = 'data'

    id = Column(Integer, primary_key=True)
    target1 = Column(PickleType)
    target2 = Column(PickleType)
    points = Column(PickleType)

    def __init__(self, target1, target2, points):
        self.target1 = target1
        self.target2 = target2
        self.points = points

Base.metadata.create_all()

class PlanetLabNode(object):
    def __init__(self, host, ip, lat, lon):
        self.host = host
        self.ip = ip
        self.lat = float(lat)
        self.lon = float(lon)
        self.connected = False

    def connect(self):
        try:
            self.connectPL()
        except Exception, e:
            print e
            try:
                self.restartPL()
                self.connectPL()
            except Exception, e:
                print e
                self.connected = False

    def connectPL(self):
        rem = SshMachine(self.host, user='ucb_268_measure', keyfile='~/.ssh/id_rsa',
                         ssh_opts=["-o StrictHostKeyChecking no",
                                   "-o UserKnownHostsFile=/dev/null"])
        conn = ssh_connect(rem, 18861, config={'allow_pickle' : True})
        assert conn.root.test() == 1, "%s: RPC Failure" % self.host
        self.conn = conn
        self.connected = True

    def restartPL(self):
        FNULL = open(os.devnull, 'w')
        subprocess.call(["ssh", "-t", "-i", "~/.ssh/id_rsa", "-o StrictHostKeyChecking no",
                         "-o UserKnownHostsFile=/dev/null", "ucb_268_measure@%s" % self.host,
                         "sudo tking-server stop; sudo tking-server start;"],
                         stdout=FNULL, stderr=FNULL)

    def handleConnExceptions(fn):
        def wrapped(self, *args, **kwargs):
            if not self.connected:
                self.connect()

            if self.connected:
                try:
                    return fn(self, *args, **kwargs)
                except Exception, e:
                    print e
                    self.connected = False
                    try:
                        self.connect()
                        return fn(self, *args, **kwargs)
                    except Exception, e:
                        print e
                        self.connected = False
                        return None
            else:
                return None
        return wrapped

    @handleConnExceptions
    def get_latency(self, target1, ip1, target2, ip2):
        timed_fn = rpyc.timed(self.conn.root.get_latency, 15)
        return timed_fn(target1, ip1, target2, ip2)

    @handleConnExceptions
    def get_k(self, target, ip):
        timed_fn = rpyc.timed(self.conn.get_k, 15)
        return timed_fn(target, ip)

    def get_distance(self, lat, lon):
        return distance((self.lat, self.lon), (lat, lon))

def select_random_points():
    target1 = open_resolvers.randomkey()
    target2 = geoip.randomkey()
    while not geoip.exists(target1):
        target1 = open_resolvers.randomkey()

    ip1, coord1 = list(all_dns.smembers(target1))[0], eval(list(geoip.smembers(target1))[0])[1:]
    ip2, coord2 = list(all_dns.smembers(target2))[0], eval(list(geoip.smembers(target2))[0])[1:]

    return (target1, ip1, coord1), (target2, ip2, coord2)

def query_latency(target1, target2):
    name1, ip1, coord1 = target1
    name2, ip2, coord2 = target2

    # Get closest 4 PL Nodes
    distances = map(lambda node: (distance(coord1, (node.lat, node.lon)), node), pl_nodes)
    distances.sort()
    distances = distances[:5]

    results = map(lambda (dist, node): (node.host, node.get_latency(name1, ip1, name2, ip2)), distances)
    return results

def one_round():
    print 'Task'
    try:
        target1, target2 = select_random_points()
        results = query_latency(target1, target2)
        return target1, target2, results
    except Exception, e:
        print e
        return None

def doWork():
    Session = sessionmaker(bind=engine)
    s = Session()
    while True:
        try:
            target1, target2, result_set = one_round()
            points = []
            if result_set:
                for host, info in filter(None,result_set):
                    if info:
                        try:
                            # TODO: Store None Responses As Well
                            (end_time, start_time, ping_times, address) = info
                            point = DataPoint(start_time, end_time, ping_times, address, host)
                            points.append(point)
                        except Exception, e:
                            print e
                            pass
                data_set = Query(target1, target2, points)
                s.add(data_set)
                s.commit()
        except Exception, e:
            print e
            pass

pl_nodes = map(lambda args: PlanetLabNode(*args), pl_hosts)

def main():
    processes = []
    for i in range(process_pool_size):
        p = Process(target=doWork)
        p.daemon = True
        p.start()
        processes.append(p)
    for p in processes:
        p.join()

main()
<commit_msg>return async result value<commit_after>import logging, os, redis, rpyc, string, subprocess
from apscheduler.scheduler import Scheduler
from datetime import datetime, timedelta
from multiprocessing import Process
from plumbum import SshMachine
from rpyc.utils.factory import ssh_connect
from utilities import distance

process_pool_size = 50

all_dns = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=0))
open_resolvers = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=1))
geoip = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=2))
pl_hosts = [line.split(' ')[0:4] for line in map(string.strip,open('pl-host-list-geo').readlines())]

from sqlalchemy import create_engine
from sqlalchemy import Column, Integer, PickleType
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base

engine = create_engine('sqlite:///latency.db', echo=False)
Base = declarative_base(bind=engine)

class DataPoint(object):
    def __init__(self, start, end, pings, address, test_point):
        self.start = start
        self.end = end
        self.pings = pings
        self.address = address
        self.test_point = test_point

class Query(Base):
    __tablename__ = 'data'

    id = Column(Integer, primary_key=True)
    target1 = Column(PickleType)
    target2 = Column(PickleType)
    points = Column(PickleType)

    def __init__(self, target1, target2, points):
        self.target1 = target1
        self.target2 = target2
        self.points = points

Base.metadata.create_all()

class PlanetLabNode(object):
    def __init__(self, host, ip, lat, lon):
        self.host = host
        self.ip = ip
        self.lat = float(lat)
        self.lon = float(lon)
        self.connected = False

    def connect(self):
        try:
            self.connectPL()
        except Exception, e:
            print e
            try:
                self.restartPL()
                self.connectPL()
            except Exception, e:
                print e
                self.connected = False

    def connectPL(self):
        rem = SshMachine(self.host, user='ucb_268_measure', keyfile='~/.ssh/id_rsa',
                         ssh_opts=["-o StrictHostKeyChecking no",
                                   "-o UserKnownHostsFile=/dev/null"])
        conn = ssh_connect(rem, 18861, config={'allow_pickle' : True})
        assert conn.root.test() == 1, "%s: RPC Failure" % self.host
        self.conn = conn
        self.connected = True

    def restartPL(self):
        FNULL = open(os.devnull, 'w')
        subprocess.call(["ssh", "-t", "-i", "~/.ssh/id_rsa", "-o StrictHostKeyChecking no",
                         "-o UserKnownHostsFile=/dev/null", "ucb_268_measure@%s" % self.host,
                         "sudo tking-server stop; sudo tking-server start;"],
                         stdout=FNULL, stderr=FNULL)

    def handleConnExceptions(fn):
        def wrapped(self, *args, **kwargs):
            if not self.connected:
                self.connect()

            if self.connected:
                try:
                    return fn(self, *args, **kwargs)
                except Exception, e:
                    print e
                    self.connected = False
                    try:
                        self.connect()
                        return fn(self, *args, **kwargs)
                    except Exception, e:
                        print e
                        self.connected = False
                        return None
            else:
                return None
        return wrapped

    @handleConnExceptions
    def get_latency(self, target1, ip1, target2, ip2):
        timed_fn = rpyc.timed(self.conn.root.get_latency, 15)
        return timed_fn(target1, ip1, target2, ip2).value()

    @handleConnExceptions
    def get_k(self, target, ip):
        timed_fn = rpyc.timed(self.conn.get_k, 15)
        return timed_fn(target, ip).value()

    def get_distance(self, lat, lon):
        return distance((self.lat, self.lon), (lat, lon))

def select_random_points():
    target1 = open_resolvers.randomkey()
    target2 = geoip.randomkey()
    while not geoip.exists(target1):
        target1 = open_resolvers.randomkey()

    ip1, coord1 = list(all_dns.smembers(target1))[0], eval(list(geoip.smembers(target1))[0])[1:]
    ip2, coord2 = list(all_dns.smembers(target2))[0], eval(list(geoip.smembers(target2))[0])[1:]

    return (target1, ip1, coord1), (target2, ip2, coord2)

def query_latency(target1, target2):
    name1, ip1, coord1 = target1
    name2, ip2, coord2 = target2

    # Get closest 4 PL Nodes
    distances = map(lambda node: (distance(coord1, (node.lat, node.lon)), node), pl_nodes)
    distances.sort()
    distances = distances[:5]

    results = map(lambda (dist, node): (node.host, node.get_latency(name1, ip1, name2, ip2)), distances)
    return results

def one_round():
    print 'Task'
    try:
        target1, target2 = select_random_points()
        results = query_latency(target1, target2)
        return target1, target2, results
    except Exception, e:
        print e
        return None

def doWork():
    Session = sessionmaker(bind=engine)
    s = Session()
    while True:
        try:
            target1, target2, result_set = one_round()
            points = []
            if result_set:
                for host, info in filter(None,result_set):
                    if info:
                        try:
                            # TODO: Store None Responses As Well
                            (end_time, start_time, ping_times, address) = info
                            point = DataPoint(start_time, end_time, ping_times, address, host)
                            points.append(point)
                        except Exception, e:
                            print e
                            pass
                data_set = Query(target1, target2, points)
                s.add(data_set)
                s.commit()
        except Exception, e:
            print e
            pass

pl_nodes = map(lambda args: PlanetLabNode(*args), pl_hosts)

def main():
    processes = []
    for i in range(process_pool_size):
        p = Process(target=doWork)
        p.daemon = True
        p.start()
        processes.append(p)
    for p in processes:
        p.join()

main()
<|endoftext|>


<commit_before>#
# Kiwi: a Framework and Enhanced Widgets for Python
#
# Copyright (C) 2005,2006 Async Open Source
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
#
# Author(s): Johan Dahlin <jdahlin@async.com.br>
#

import os
import gettext

import atk
import gtk

__all__ = ['error', 'info', 'messagedialog', 'warning', 'yesno', 'save', 'open',
           'HIGAlertDialog', 'BaseDialog']

_ = gettext.gettext

_IMAGE_TYPES = {
    gtk.MESSAGE_INFO: gtk.STOCK_DIALOG_INFO,
    gtk.MESSAGE_WARNING : gtk.STOCK_DIALOG_WARNING,
    gtk.MESSAGE_QUESTION : gtk.STOCK_DIALOG_QUESTION,
    gtk.MESSAGE_ERROR : gtk.STOCK_DIALOG_ERROR,
}

_BUTTON_TYPES = {
    gtk.BUTTONS_NONE: (),
    gtk.BUTTONS_OK: (gtk.STOCK_OK, gtk.RESPONSE_OK,),
    gtk.BUTTONS_CLOSE: (gtk.STOCK_CLOSE, gtk.RESPONSE_CLOSE,),
    gtk.BUTTONS_CANCEL: (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL,),
    gtk.BUTTONS_YES_NO: (gtk.STOCK_NO, gtk.RESPONSE_NO,
                         gtk.STOCK_YES, gtk.RESPONSE_YES),
    gtk.BUTTONS_OK_CANCEL: (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL,
                            gtk.STOCK_OK, gtk.RESPONSE_OK)
    }

class HIGAlertDialog(gtk.Dialog):
    def __init__(self, parent, flags,
                 type=gtk.MESSAGE_INFO, buttons=gtk.BUTTONS_NONE):
        if not type in _IMAGE_TYPES:
            raise TypeError(
                "type must be one of: %s", ', '.join(_IMAGE_TYPES.keys()))
        if not buttons in _BUTTON_TYPES:
            raise TypeError(
                "buttons be one of: %s", ', '.join(_BUTTON_TYPES.keys()))

        gtk.Dialog.__init__(self, '', parent, flags)
	self.set_border_width(5)
	self.set_resizable(False)
	self.set_has_separator(False)
	self.set_title("")
	self.set_skip_taskbar_hint(True)
	self.vbox.set_spacing(14)
	self.get_accessible().set_role(atk.ROLE_ALERT)

	self._primary_label = gtk.Label()
	self._secondary_label = gtk.Label()
	self._details_label = gtk.Label()
	self._image = gtk.image_new_from_stock(_IMAGE_TYPES[type],
                                               gtk.ICON_SIZE_DIALOG)
	self._image.set_alignment(0.5, 0.0)

	self._primary_label.set_use_markup(True)
        for label in (self._primary_label, self._secondary_label,
                      self._details_label):
            label.set_line_wrap(True)
            label.set_selectable(True)
            label.set_alignment(0.0, 0.5)

        hbox = gtk.HBox(False, 12)
        hbox.set_border_width(5)
        hbox.pack_start(self._image, False, False)

	vbox = gtk.VBox(False, 0)
        hbox.pack_start(vbox, False, False)
        vbox.pack_start(self._primary_label, False, False)
        vbox.pack_start(self._secondary_label, False, False)

	self._expander = gtk.expander_new_with_mnemonic(_("Show more _details"))
        self._expander.set_spacing(6)
        self._expander.add(self._details_label)
        vbox.pack_start(self._expander, False, False)
        self.vbox.pack_start(hbox, False, False)
        hbox.show_all()
        self._expander.hide()
        self.add_buttons(*_BUTTON_TYPES[buttons])

    def set_primary(self, text):
        self._primary_label.set_markup("%s" % text)

    def set_secondary(self, text):
        self._secondary_label.set_markup(text)

    def set_details(self, text):
        self._details_label.set_text(text)
        self._expander.show()

    def set_details_widget(self, widget):
        self._expander.remove(self._details_label)
        self._expander.add(widget)
        widget.show()
        self._expander.show()

class BaseDialog(gtk.Dialog):
    def __init__(self, parent=None, title='', flags=0, buttons=()):
        if parent and not isinstance(parent, gtk.Window):
            raise TypeError("parent needs to be None or a gtk.Window subclass")

        if not flags and parent:
            flags &= (gtk.DIALOG_MODAL |
                      gtk.DIALOG_DESTROY_WITH_PARENT)

        gtk.Dialog.__init__(self, title=title, parent=parent,
                            flags=flags, buttons=buttons)

def messagedialog(dialog_type, short, long=None, parent=None,
                  buttons=gtk.BUTTONS_OK, default=-1):
    """Create and show a MessageDialog.

    @param dialog_type: one of constants
      - gtk.MESSAGE_INFO
      - gtk.MESSAGE_WARNING
      - gtk.MESSAGE_QUESTION
      - gtk.MESSAGE_ERROR
    @param short:       A header text to be inserted in the dialog.
    @param long:        A long description of message.
    @param parent:      The parent widget of this dialog
    @type parent:       a gtk.Window subclass
    @param buttons:     The button type that the dialog will be display,
      one of the constants:
       - gtk.BUTTONS_NONE
       - gtk.BUTTONS_OK
       - gtk.BUTTONS_CLOSE
       - gtk.BUTTONS_CANCEL
       - gtk.BUTTONS_YES_NO
       - gtk.BUTTONS_OK_CANCEL
      or a tuple or 2-sized tuples representing label and response. If label
      is a stock-id a stock icon will be displayed.
    @param default: optional default response id
    """
    if buttons in (gtk.BUTTONS_NONE, gtk.BUTTONS_OK, gtk.BUTTONS_CLOSE,
                   gtk.BUTTONS_CANCEL, gtk.BUTTONS_YES_NO,
                   gtk.BUTTONS_OK_CANCEL):
        dialog_buttons = buttons
        buttons = []
    else:
        if type(buttons) != tuple:
            raise TypeError(
                "buttons must be a GtkButtonsTypes constant or a tuple")
        dialog_buttons = gtk.BUTTONS_NONE

    if parent and not isinstance(parent, gtk.Window):
        raise TypeError("parent must be a gtk.Window subclass")

    d = HIGAlertDialog(parent=parent, flags=gtk.DIALOG_MODAL,
                       type=dialog_type, buttons=dialog_buttons)
    for text, response in buttons:
        d.add_buttons(text, response)

    d.set_primary(short)

    if long:
        if isinstance(long, gtk.Widget):
            d.set_details_widget(long)
        elif isinstance(long, basestring):
            d.set_details(long)
        else:
            raise TypeError("long must be a gtk.Widget or a string")

    if default != -1:
        d.set_default_response(default)

    if parent:
        d.set_transient_for(parent)
        d.set_modal(True)

    response = d.run()
    d.destroy()
    return response

def _simple(type, short, long=None, parent=None, buttons=gtk.BUTTONS_OK,
          default=-1):
    if buttons == gtk.BUTTONS_OK:
        default = gtk.RESPONSE_OK
    return messagedialog(type, short, long,
                         parent=parent, buttons=buttons,
                         default=default)

def error(short, long=None, parent=None, buttons=gtk.BUTTONS_OK, default=-1):
    return _simple(gtk.MESSAGE_ERROR, short, long, parent=parent,
                   buttons=buttons, default=default)

def info(short, long=None, parent=None, buttons=gtk.BUTTONS_OK, default=-1):
    return _simple(gtk.MESSAGE_INFO, short, long, parent=parent,
                   buttons=buttons, default=default)

def warning(short, long=None, parent=None, buttons=gtk.BUTTONS_OK, default=-1):
    return _simple(gtk.MESSAGE_WARNING, short, long, parent=parent,
                   buttons=buttons, default=default)

def yesno(text, parent=None, default=gtk.RESPONSE_YES):
    return messagedialog(gtk.MESSAGE_WARNING, text, None, parent,
                         buttons=gtk.BUTTONS_YES_NO,
                         default=default)

def open(title='', parent=None, patterns=[], folder=None):
    """Displays an open dialog."""
    filechooser = gtk.FileChooserDialog(title or _('Open'),
                                        parent,
                                        gtk.FILE_CHOOSER_ACTION_OPEN,
                                        (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL,
                                         gtk.STOCK_OPEN, gtk.RESPONSE_OK))

    if patterns:
        file_filter = gtk.FileFilter()
        for pattern in patterns:
            file_filter.add_pattern(pattern)
        filechooser.set_filter(file_filter)
    filechooser.set_default_response(gtk.RESPONSE_OK)

    if folder:
        filechooser.set_current_folder(folder)

    response = filechooser.run()
    if response != gtk.RESPONSE_OK:
        filechooser.destroy()
        return

    path = filechooser.get_filename()
    if path and os.access(path, os.R_OK):
        filechooser.destroy()
        return path

    abspath = os.path.abspath(path)

    error(_('Could not open file "%s"') % abspath,
          _('The file "%s" could not be opened. '
            'Permission denied.') %  abspath)

    filechooser.destroy()
    return path

def save(title='', parent=None, current_name='', folder=None):
    """Displays a save dialog."""
    filechooser = gtk.FileChooserDialog(title or _('Save'),
                                        parent,
                                        gtk.FILE_CHOOSER_ACTION_SAVE,
                                        (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL,
                                         gtk.STOCK_SAVE, gtk.RESPONSE_OK))
    if current_name:
        filechooser.set_current_name(current_name)
    filechooser.set_default_response(gtk.RESPONSE_OK)

    if folder:
        filechooser.set_current_folder(folder)

    path = None
    while True:
        response = filechooser.run()
        if response != gtk.RESPONSE_OK:
            path = None
            break

        path = filechooser.get_filename()
        if not os.path.exists(path):
            break

        submsg1 = _('A file named "%s" already exists') % os.path.abspath(path)
        submsg2 = _('Do you wish to replace it with the current project?')
        text = '%s\n\n%s\n' % \
              (submsg1, submsg2)
        result = messagedialog(gtk.MESSAGE_ERROR,
                               text,
                               parent=parent,
                               buttons=(gtk.STOCK_CANCEL,
                                        gtk.RESPONSE_CANCEL,
                                        _("Replace"),
                                        gtk.RESPONSE_YES))
        # the user want to overwrite the file
        if result == gtk.RESPONSE_YES:
            break

    filechooser.destroy()
    return path

def _test():
     yesno('Kill?', default=gtk.RESPONSE_NO)

     info('Some information displayed not too long\nbut not too short',
          long=('foobar ba asdjaiosjd oiadjoisjaoi aksjdasdasd kajsdhakjsdh\n'
                'askdjhaskjdha skjdhasdasdjkasldj alksdjalksjda lksdjalksdj\n'
                'asdjaslkdj alksdj lkasjdlkjasldkj alksjdlkasjdlkasjdlka jklsdjakls\n'
                'ask;ldjaklsjdlkasjd alksdj laksjdlkasjd lkajs lkjdl kjaslk jkl\n'),
          default=gtk.RESPONSE_OK,
          )

     error('An error occurred', gtk.Button('Woho'))
     error('Unable to mount the selected volume.',
           'mount: can\'t find /media/cdrom0 in /etc/fstab or /etc/mtab')
     print open(title='Open a file', patterns=['*.py'])
     print save(title='Save a file', current_name='foobar.py')

if __name__ == '__main__':
    _test()
<commit_msg>Fixing a bug -- buttons and its response codes must be specified inside tuples. r=jdahlin<commit_after>#
# Kiwi: a Framework and Enhanced Widgets for Python
#
# Copyright (C) 2005,2006 Async Open Source
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
#
# Author(s): Johan Dahlin <jdahlin@async.com.br>
#

import os
import gettext

import atk
import gtk

__all__ = ['error', 'info', 'messagedialog', 'warning', 'yesno', 'save', 'open',
           'HIGAlertDialog', 'BaseDialog']

_ = gettext.gettext

_IMAGE_TYPES = {
    gtk.MESSAGE_INFO: gtk.STOCK_DIALOG_INFO,
    gtk.MESSAGE_WARNING : gtk.STOCK_DIALOG_WARNING,
    gtk.MESSAGE_QUESTION : gtk.STOCK_DIALOG_QUESTION,
    gtk.MESSAGE_ERROR : gtk.STOCK_DIALOG_ERROR,
}

_BUTTON_TYPES = {
    gtk.BUTTONS_NONE: (),
    gtk.BUTTONS_OK: (gtk.STOCK_OK, gtk.RESPONSE_OK,),
    gtk.BUTTONS_CLOSE: (gtk.STOCK_CLOSE, gtk.RESPONSE_CLOSE,),
    gtk.BUTTONS_CANCEL: (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL,),
    gtk.BUTTONS_YES_NO: (gtk.STOCK_NO, gtk.RESPONSE_NO,
                         gtk.STOCK_YES, gtk.RESPONSE_YES),
    gtk.BUTTONS_OK_CANCEL: (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL,
                            gtk.STOCK_OK, gtk.RESPONSE_OK)
    }

class HIGAlertDialog(gtk.Dialog):
    def __init__(self, parent, flags,
                 type=gtk.MESSAGE_INFO, buttons=gtk.BUTTONS_NONE):
        if not type in _IMAGE_TYPES:
            raise TypeError(
                "type must be one of: %s", ', '.join(_IMAGE_TYPES.keys()))
        if not buttons in _BUTTON_TYPES:
            raise TypeError(
                "buttons be one of: %s", ', '.join(_BUTTON_TYPES.keys()))

        gtk.Dialog.__init__(self, '', parent, flags)
	self.set_border_width(5)
	self.set_resizable(False)
	self.set_has_separator(False)
	self.set_title("")
	self.set_skip_taskbar_hint(True)
	self.vbox.set_spacing(14)
	self.get_accessible().set_role(atk.ROLE_ALERT)

	self._primary_label = gtk.Label()
	self._secondary_label = gtk.Label()
	self._details_label = gtk.Label()
	self._image = gtk.image_new_from_stock(_IMAGE_TYPES[type],
                                               gtk.ICON_SIZE_DIALOG)
	self._image.set_alignment(0.5, 0.0)

	self._primary_label.set_use_markup(True)
        for label in (self._primary_label, self._secondary_label,
                      self._details_label):
            label.set_line_wrap(True)
            label.set_selectable(True)
            label.set_alignment(0.0, 0.5)

        hbox = gtk.HBox(False, 12)
        hbox.set_border_width(5)
        hbox.pack_start(self._image, False, False)

	vbox = gtk.VBox(False, 0)
        hbox.pack_start(vbox, False, False)
        vbox.pack_start(self._primary_label, False, False)
        vbox.pack_start(self._secondary_label, False, False)

	self._expander = gtk.expander_new_with_mnemonic(_("Show more _details"))
        self._expander.set_spacing(6)
        self._expander.add(self._details_label)
        vbox.pack_start(self._expander, False, False)
        self.vbox.pack_start(hbox, False, False)
        hbox.show_all()
        self._expander.hide()
        self.add_buttons(*_BUTTON_TYPES[buttons])

    def set_primary(self, text):
        self._primary_label.set_markup("%s" % text)

    def set_secondary(self, text):
        self._secondary_label.set_markup(text)

    def set_details(self, text):
        self._details_label.set_text(text)
        self._expander.show()

    def set_details_widget(self, widget):
        self._expander.remove(self._details_label)
        self._expander.add(widget)
        widget.show()
        self._expander.show()

class BaseDialog(gtk.Dialog):
    def __init__(self, parent=None, title='', flags=0, buttons=()):
        if parent and not isinstance(parent, gtk.Window):
            raise TypeError("parent needs to be None or a gtk.Window subclass")

        if not flags and parent:
            flags &= (gtk.DIALOG_MODAL |
                      gtk.DIALOG_DESTROY_WITH_PARENT)

        gtk.Dialog.__init__(self, title=title, parent=parent,
                            flags=flags, buttons=buttons)

def messagedialog(dialog_type, short, long=None, parent=None,
                  buttons=gtk.BUTTONS_OK, default=-1):
    """Create and show a MessageDialog.

    @param dialog_type: one of constants
      - gtk.MESSAGE_INFO
      - gtk.MESSAGE_WARNING
      - gtk.MESSAGE_QUESTION
      - gtk.MESSAGE_ERROR
    @param short:       A header text to be inserted in the dialog.
    @param long:        A long description of message.
    @param parent:      The parent widget of this dialog
    @type parent:       a gtk.Window subclass
    @param buttons:     The button type that the dialog will be display,
      one of the constants:
       - gtk.BUTTONS_NONE
       - gtk.BUTTONS_OK
       - gtk.BUTTONS_CLOSE
       - gtk.BUTTONS_CANCEL
       - gtk.BUTTONS_YES_NO
       - gtk.BUTTONS_OK_CANCEL
      or a tuple or 2-sized tuples representing label and response. If label
      is a stock-id a stock icon will be displayed.
    @param default: optional default response id
    """
    if buttons in (gtk.BUTTONS_NONE, gtk.BUTTONS_OK, gtk.BUTTONS_CLOSE,
                   gtk.BUTTONS_CANCEL, gtk.BUTTONS_YES_NO,
                   gtk.BUTTONS_OK_CANCEL):
        dialog_buttons = buttons
        buttons = []
    else:
        if type(buttons) != tuple:
            raise TypeError(
                "buttons must be a GtkButtonsTypes constant or a tuple")
        dialog_buttons = gtk.BUTTONS_NONE

    if parent and not isinstance(parent, gtk.Window):
        raise TypeError("parent must be a gtk.Window subclass")

    d = HIGAlertDialog(parent=parent, flags=gtk.DIALOG_MODAL,
                       type=dialog_type, buttons=dialog_buttons)
    for text, response in buttons:
        d.add_buttons(text, response)

    d.set_primary(short)

    if long:
        if isinstance(long, gtk.Widget):
            d.set_details_widget(long)
        elif isinstance(long, basestring):
            d.set_details(long)
        else:
            raise TypeError("long must be a gtk.Widget or a string")

    if default != -1:
        d.set_default_response(default)

    if parent:
        d.set_transient_for(parent)
        d.set_modal(True)

    response = d.run()
    d.destroy()
    return response

def _simple(type, short, long=None, parent=None, buttons=gtk.BUTTONS_OK,
          default=-1):
    if buttons == gtk.BUTTONS_OK:
        default = gtk.RESPONSE_OK
    return messagedialog(type, short, long,
                         parent=parent, buttons=buttons,
                         default=default)

def error(short, long=None, parent=None, buttons=gtk.BUTTONS_OK, default=-1):
    return _simple(gtk.MESSAGE_ERROR, short, long, parent=parent,
                   buttons=buttons, default=default)

def info(short, long=None, parent=None, buttons=gtk.BUTTONS_OK, default=-1):
    return _simple(gtk.MESSAGE_INFO, short, long, parent=parent,
                   buttons=buttons, default=default)

def warning(short, long=None, parent=None, buttons=gtk.BUTTONS_OK, default=-1):
    return _simple(gtk.MESSAGE_WARNING, short, long, parent=parent,
                   buttons=buttons, default=default)

def yesno(text, parent=None, default=gtk.RESPONSE_YES):
    return messagedialog(gtk.MESSAGE_WARNING, text, None, parent,
                         buttons=gtk.BUTTONS_YES_NO,
                         default=default)

def open(title='', parent=None, patterns=[], folder=None):
    """Displays an open dialog."""
    filechooser = gtk.FileChooserDialog(title or _('Open'),
                                        parent,
                                        gtk.FILE_CHOOSER_ACTION_OPEN,
                                        (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL,
                                         gtk.STOCK_OPEN, gtk.RESPONSE_OK))

    if patterns:
        file_filter = gtk.FileFilter()
        for pattern in patterns:
            file_filter.add_pattern(pattern)
        filechooser.set_filter(file_filter)
    filechooser.set_default_response(gtk.RESPONSE_OK)

    if folder:
        filechooser.set_current_folder(folder)

    response = filechooser.run()
    if response != gtk.RESPONSE_OK:
        filechooser.destroy()
        return

    path = filechooser.get_filename()
    if path and os.access(path, os.R_OK):
        filechooser.destroy()
        return path

    abspath = os.path.abspath(path)

    error(_('Could not open file "%s"') % abspath,
          _('The file "%s" could not be opened. '
            'Permission denied.') %  abspath)

    filechooser.destroy()
    return path

def save(title='', parent=None, current_name='', folder=None):
    """Displays a save dialog."""
    filechooser = gtk.FileChooserDialog(title or _('Save'),
                                        parent,
                                        gtk.FILE_CHOOSER_ACTION_SAVE,
                                        (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL,
                                         gtk.STOCK_SAVE, gtk.RESPONSE_OK))
    if current_name:
        filechooser.set_current_name(current_name)
    filechooser.set_default_response(gtk.RESPONSE_OK)

    if folder:
        filechooser.set_current_folder(folder)

    path = None
    while True:
        response = filechooser.run()
        if response != gtk.RESPONSE_OK:
            path = None
            break

        path = filechooser.get_filename()
        if not os.path.exists(path):
            break

        submsg1 = _('A file named "%s" already exists') % os.path.abspath(path)
        submsg2 = _('Do you wish to replace it with the current project?')
        text = '%s\n\n%s\n' % \
              (submsg1, submsg2)
        result = messagedialog(gtk.MESSAGE_ERROR,
                               text,
                               parent=parent,
                               buttons=((gtk.STOCK_CANCEL,
                                        gtk.RESPONSE_CANCEL),
                                        (_("Replace"),
                                         gtk.RESPONSE_YES)))
        # the user want to overwrite the file
        if result == gtk.RESPONSE_YES:
            break

    filechooser.destroy()
    return path

def _test():
     yesno('Kill?', default=gtk.RESPONSE_NO)

     info('Some information displayed not too long\nbut not too short',
          long=('foobar ba asdjaiosjd oiadjoisjaoi aksjdasdasd kajsdhakjsdh\n'
                'askdjhaskjdha skjdhasdasdjkasldj alksdjalksjda lksdjalksdj\n'
                'asdjaslkdj alksdj lkasjdlkjasldkj alksjdlkasjdlkasjdlka jklsdjakls\n'
                'ask;ldjaklsjdlkasjd alksdj laksjdlkasjd lkajs lkjdl kjaslk jkl\n'),
          default=gtk.RESPONSE_OK,
          )

     error('An error occurred', gtk.Button('Woho'))
     error('Unable to mount the selected volume.',
           'mount: can\'t find /media/cdrom0 in /etc/fstab or /etc/mtab')
     print open(title='Open a file', patterns=['*.py'])
     print save(title='Save a file', current_name='foobar.py')

if __name__ == '__main__':
    _test()
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A python script to normalize a wide variety of value notations
from electronics engineering.

Examples of valid notations include:
    1,234.56kΩ
    1k234
    1k234Ω
    1,234.56Ω
    4µA
    4e6A
    4e6nA

Usage example:
    >>> print(normalizeEngineerInput("1µ234 Ω"))
    (1.234e-6, 'Ω')

Originally published at techoverflow.net.
"""
import re
import math
import itertools
import functools
from toolz import functoolz
from collections import namedtuple
import numpy as np

__all__ = ["Quantity", "UnannotatedReturnValueError",
           "normalize_interpunctation", "EngineerIO"]

Quantity = namedtuple("Quantity", ["unit"])

class UnannotatedReturnValueError(Exception):
    pass

# Suffices handled by the library
_default_suffices = [["y"], ["z"], ["a"], ["f"], ["p"], ["n"], ["µ", "u"], ["m"], [],
                     ["k"], ["M"], ["G"], ["T"], ["E"], ["Z"], ["Y"]]
_default_1st_suffix_exp = -24  # The exponential multiplier for the first suffix

# Valid unit designators. Ensure no SI suffix is added here
_default_units = frozenset(['F', 'A', 'Ω', 'W', 'H', 'C', 'K', 'Hz', 'V', 'J', 'S'])
_numeric_allowed = set("0123456789-e.")

class EngineerIO(object):
    instance = None
    """Default instance, used for global functions. Initialized on first use"""
    def __init__(self, units=_default_units,
                 unit_prefixes="Δ°",
                 suffices=_default_suffices,
                 first_suffix_exp=_default_1st_suffix_exp):
        """
        Initialize a new EngineerIO instance with default or custom suffix

        Parameters:
        -----------
        units : iterable of strings
            An iterable of valid units (1-char or 2-char)
        unit_prefixes : string
            A list of prefixes that are silently ignored.
            Constraint: unitPrefixes ∩ suffices == ∅
        suffices : list of lists of unit strings
            For each SI exponent, a list of valid suffix strings for each exponent.
            Each successive element in the list is 1e3 from the previous one.
            For generating strings from numbers, the first suffix in each nested list is preferred
            The list must not contain the empty string but shall be empty if no suffix shall be used
        first_suffix_exp : int
            The decimal exponent of the first suffix in the suffix list.
        """
        self.units = set(units)
        self.suffices = suffices
        self.first_suffix_exp = first_suffix_exp
        # Unit prefixes will only be used in strip, so we can strip spaces in one go.
        self.strippable = unit_prefixes + " \t\n"
        # Compute maps
        self.all_suffixes = set(itertools.chain(*self.suffices))
        self._recompute_suffix_maps()

    def _recompute_suffix_maps(self):
        """
        Recompute the exponent -> suffix map and
        the suffix -> exponent map
        """
        # Compute inverse suffix map
        self.exp_suffix_map = {}  # Key: exp // 3, Value: suffix
        self.suffix_exp_map = {}  # Key: suffix, value: exponent
        current_exp = self.first_suffix_exp
        # Iterate over first suffices in each list
        for current_suffices in self.suffices:
            if not current_suffices: # No suffix to be used
                # Do not self suffix exp map
                self.exp_suffix_map[current_exp // 3] = ""
                current_exp += 3
                continue
            # Compute exponent -> suffix (only first suffix)
            self.exp_suffix_map[current_exp // 3] = current_suffices[0]
            # Compute suffix -> exponent
            for current_suffix in current_suffices:
                self.suffix_exp_map[current_suffix] = current_exp
            current_exp += 3
        # Compute min/max SI value
        self.exp_map_min = min(self.exp_suffix_map.keys())
        self.exp_map_max = max(self.exp_suffix_map.keys())

    def split_input(self, s):
        """
        Separate a string into a 3-tuple (number, suffix, unit).
        returns None if the string could not be parsed.

        The tuple will never contain None but empty strings if some
        element is not present. The number must be present for the string
        to be considered valid.

        Units always need to be a suffix. Instead
        Thousands separators or commata instead of points may be used
        (see normalize_interpunctation documentation).

        Thousands separators and suffix-as-decimal-separators may NOT
        be mixed. Whitespace is removed automatically.
        """
        # Remove thousands separator & ensure dot is used
        s = normalize_interpunctation(s).replace(" ", "")
        s, unit = self.split_unit(s) # Remove unit
        # Try to find SI suffix at the end or in the middle
        if s[-1] in self.all_suffixes:
            s, suffix = s[:-1], s[-1]
        else:  # Try to find unit anywhere
            isSuffixList = [ch in self.all_suffixes for ch in s]
            # Ensure only ONE unit occurs in the string
            suffixCount = sum(isSuffixList)
            if suffixCount > 1:
                raise ValueError("More than one SI suffix in the string")
            elif suffixCount == 0:
                suffix = ""
            else:  # suffixCount == 1 => correct
                # Suffix-as-decimal-separator --> no other decimal separator
                if "." in s:  # Commata already handled by normalize_interpunctation
                    return None
                suffixIndex = isSuffixList.index(True)
                # Suffix must NOT be first character
                if suffixIndex == 0:
                    raise ValueError("Suffix in '{0}' must not be the first char".format(s))
                suffix = s[suffixIndex]
                s = s.replace(suffix, ".")
        # Handle unit prefix (if any). Not allowable if no unit is present
        s = s.strip(self.strippable)
        # Final check: Is there any number left and is it valid?
        if not all((ch in _numeric_allowed for ch in s)):
            raise ValueError("Remainder of string is not purely numeric: {0}".format(s))
        return (s, suffix, unit)

    def split_unit(self, s):
        """
        Split a string into (remainder, unit).
        Only units in the units set are recognized
        unit may be "None" if no unit is recognized
        """
        # Handle 2-character units (e.g. 'Hz'), then 1-character unity (e.g. 'V')
        if s[-2:] in self.units: # Will also handle unit-only 1-char strings
            s, unit = s[:-2], s[-2:]
        elif s[-1] in self.units:  # Handle 1-char units
            s, unit = s[:-1], s[-1]
        else: # No unit
            s, unit = s, ''
        # Remove unit prefix, if any (e.g. degrees symbol, delta symbol)
        s = s.rstrip(self.strippable)
        return s, unit

    def normalize(self, s, encoding="utf8"):
        """
        Converts an engineer's input of a wide variety of formats to a numeric
        value.

        Returns a pair (number, unit) or None if the conversion could not be performed.

        See splitSuffixSeparator() for further details on supported formats
        """
        # Make sure it's a decoded string
        if isinstance(s, bytes):
            s = s.decode(encoding)
        # Perform splitting
        res = self.split_input(s.strip())
        (num, suffix, unit) = res
        mul = (10 ** self.suffix_exp_map[suffix]) if suffix else 1
        return (float(num) * mul, unit)

    def safe_normalize(self, s, encoding="utf8"):
        """
        Same as normalize(), but returns None instead of raising
        on error.
        """
        try:
            return self.normalize(s, encoding)
        except ValueError:
            return None

    def format(self, v, unit=""):
        """
        Format v using SI suffices with optional units.
        Produces a string with 3 visible digits.
        """
        #Suffix map is indexed by one third of the decadic logarithm.
        exp = 0 if v == 0. else math.log(abs(v), 10.)
        suffixMapIdx = int(math.floor(exp / 3.))
        #Ensure we're in range
        if not self.exp_map_min < suffixMapIdx < self.exp_map_max:
            raise ValueError("Value out of range: {0}".format(v))
        #Pre-multiply the value
        v = v * (10.0 ** -(suffixMapIdx * 3))
        #Delegate the rest of the task to the helper
        return _formatWithSuffix(v, self.exp_suffix_map[suffixMapIdx] + unit)

    def auto_suffix_1d(self, arr):
        """
        Takes an array of arbitrary values and determines
        what is the best suffix (e.g. M, m, n, f) to represent
        as many values as possible with as few powers of 10 as possible.

        Returns a tuple (factor, suffix) where the factor is a floating-point
        value to multiply the array with to obtain value with "suffix" suffix.
        """
        # Compute logarithmic magnitudes of data
        arr_log = np.log10(np.abs(arr))
        arr_log[np.isinf(arr_log)] = 0  # log(0) == inf
        log_mean = arr_log.mean()
        # Generate score matrix
        suffix_idx = int(round(log_mean / 3))
        # Ensure we're in range
        suffix_idx = max(self.exp_map_min, suffix_idx)
        suffix_idx = min(self.exp_map_max, suffix_idx)
        # Pre-multiply the value
        multiplier = 10.0 ** -(suffix_idx * 3)
        return multiplier, self.exp_suffix_map[suffix_idx]

    def auto_format(self, fn, *args, **kwargs):
        """
        Auto-format a value by leveraging function annotations.
        For example this can be used to format using UliEngineering Physics quantities
        The function's return value is expected to a be annotated with a Quantity() value.
        """
        if not callable(fn):
            raise ValueError("fn must be callable")
        # Access innermost function inside possibly nested partials
        annotatedFN = fn
        while isinstance(annotatedFN, functools.partial):
            annotatedFN = annotatedFN.func
        try:
            qty = annotatedFN.__annotations__["return"]
        except KeyError:
            raise UnannotatedReturnValueError("Function {0} does not have an annotated return value")
        return self.format(fn(*args, **kwargs), qty.unit)

    def normalize_numeric_safe(self, arg):
        """
        Normalize each element of an iterable and retrieve only the numeric value
        (the unit is ignored). Works on iterables and string-likes.

        Use toolz.itertoolz.compact() on the result to remove all None values.

        Returns a list with None (on error) or the numeric value (no unit)
        """
        if isinstance(arg, int) or isinstance(arg, float):
            return arg
        # If it's stringlike, apply directly
        if isinstance(arg, str) or isinstance(arg, bytes):
            v = self.safe_normalize(arg)
            return None if v is None else v[0]
        # It's an iterable
        ret = []
        for elem in arg:
            v = self.safe_normalize(elem)
            ret.append(None if v is None else v[0])
        return ret

    def normalize_numeric(self, arg):
        """
        Normalize each element of an iterable and retrieve only the numeric value
        (the unit is ignored). Works on iterables and string-likes.

        Raises if any of the values can't be normalized
        """
        if isinstance(arg, int) or isinstance(arg, float):
            return arg
        # If it's stringlike, apply directly
        if isinstance(arg, str) or isinstance(arg, bytes):
            return self.normalize(arg)[0]
        # It's an iterable
        return [self.normalize(elem)[0] for elem in arg]

# Default instance, initialized on first use

def __init_engineer_io_instance():
    """
    Helper function to initialize the default instance
    the first time any of the global functions is used.
    """
    if EngineerIO.instance is None:
        EngineerIO.instance = EngineerIO()

__replace_comma_dot = lambda s: s.replace(",", ".")
"""
Map of a transform to apply to the string
during interpunctation normalization,
depending on (commaFound, dotFound, commaFoundFirst).
Must contain every possible variant
"""
_interpunct_transform_map = {
    # Found nothing or only point -> no modification required
    (False, False, False): functoolz.identity,
    (False, False, True): functoolz.identity,
    (False, True, False): functoolz.identity,
    (False, True, True): functoolz.identity,
    # Only comma -> replace and exit
    (True, False, False): __replace_comma_dot,
    (True, False, True): __replace_comma_dot,
    # Below this line: Both comma and dot found
    # Comma first => comma used as thousands separators
    (True, True, True): lambda s: s.replace(",", ""),
    # Dot first => dot used as thousands separator
    (True, True, False): lambda s: s.replace(".", "").replace(",", ".")
}

def normalize_interpunctation(s):
    """
    Normalize comma to point for float conversion.
    Correctly handles thousands separators.

    Note that cases like "1,234" are undecidable between
    "1234" and "1.234". They are treated as "1.234".

    Only points and commata are potentially modified.
    Other characters and digits are not handled.
    """
    foundCommaFirst = False
    commaIdx = s.find(",")
    pointIdx = s.find(".")
    foundComma = commaIdx is not None
    foundPoint = pointIdx is not None
    commaFirst = commaIdx < pointIdx if (foundComma and foundPoint) else None
    # Get the appropriate transform function from the map an run it on s
    return _interpunct_transform_map[(foundComma, foundPoint, commaFirst)](s)

def _formatWithSuffix(v, suffix=""):
    """
    Format a given value with a given suffix.
    This helper function formats the value to 3 visible digits.
    v must be pre-multiplied by the factor implied by the suffix
    """
    if v < 10.0:
        res = "{:.2f}".format(v)
    elif v < 100.0:
        res = "{:.1f}".format(v)
    else:  # Should only happen if v < 1000
        res = str(int(round(v)))
    #Avoid appending whitespace if there is no suffix
    return "{0} {1}".format(res, suffix) if suffix else res

def normalizeEngineerInput(s, encoding="utf8"):
    __init_engineer_io_instance()
    return EngineerIO.instance.normalize(s)

def formatValue(v, unit=""):
    __init_engineer_io_instance()
    return EngineerIO.instance.normalize(v, unit)

def safeNormalizeEngineerInput(v, unit=""):
    __init_engineer_io_instance()
    return EngineerIO.instance.safe_normalize(v, unit)

def normalize_numeric(v):
    __init_engineer_io_instance()
    return EngineerIO.instance.normalize_numeric(v)

def auto_format(v, *args, **kwargs):
    __init_engineer_io_instance()
    return EngineerIO.instance.auto_format(v, *args, **kwargs)
<commit_msg>Add split_unit fallback for too short strings<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A python script to normalize a wide variety of value notations
from electronics engineering.

Examples of valid notations include:
    1,234.56kΩ
    1k234
    1k234Ω
    1,234.56Ω
    4µA
    4e6A
    4e6nA

Usage example:
    >>> print(normalizeEngineerInput("1µ234 Ω"))
    (1.234e-6, 'Ω')

Originally published at techoverflow.net.
"""
import re
import math
import itertools
import functools
from toolz import functoolz
from collections import namedtuple
import numpy as np

__all__ = ["Quantity", "UnannotatedReturnValueError",
           "normalize_interpunctation", "EngineerIO"]

Quantity = namedtuple("Quantity", ["unit"])

class UnannotatedReturnValueError(Exception):
    pass

# Suffices handled by the library
_default_suffices = [["y"], ["z"], ["a"], ["f"], ["p"], ["n"], ["µ", "u"], ["m"], [],
                     ["k"], ["M"], ["G"], ["T"], ["E"], ["Z"], ["Y"]]
_default_1st_suffix_exp = -24  # The exponential multiplier for the first suffix

# Valid unit designators. Ensure no SI suffix is added here
_default_units = frozenset(['F', 'A', 'Ω', 'W', 'H', 'C', 'K', 'Hz', 'V', 'J', 'S'])
_numeric_allowed = set("0123456789-e.")

class EngineerIO(object):
    instance = None
    """Default instance, used for global functions. Initialized on first use"""
    def __init__(self, units=_default_units,
                 unit_prefixes="Δ°",
                 suffices=_default_suffices,
                 first_suffix_exp=_default_1st_suffix_exp):
        """
        Initialize a new EngineerIO instance with default or custom suffix

        Parameters:
        -----------
        units : iterable of strings
            An iterable of valid units (1-char or 2-char)
        unit_prefixes : string
            A list of prefixes that are silently ignored.
            Constraint: unitPrefixes ∩ suffices == ∅
        suffices : list of lists of unit strings
            For each SI exponent, a list of valid suffix strings for each exponent.
            Each successive element in the list is 1e3 from the previous one.
            For generating strings from numbers, the first suffix in each nested list is preferred
            The list must not contain the empty string but shall be empty if no suffix shall be used
        first_suffix_exp : int
            The decimal exponent of the first suffix in the suffix list.
        """
        self.units = set(units)
        self.suffices = suffices
        self.first_suffix_exp = first_suffix_exp
        # Unit prefixes will only be used in strip, so we can strip spaces in one go.
        self.strippable = unit_prefixes + " \t\n"
        # Compute maps
        self.all_suffixes = set(itertools.chain(*self.suffices))
        self._recompute_suffix_maps()

    def _recompute_suffix_maps(self):
        """
        Recompute the exponent -> suffix map and
        the suffix -> exponent map
        """
        # Compute inverse suffix map
        self.exp_suffix_map = {}  # Key: exp // 3, Value: suffix
        self.suffix_exp_map = {}  # Key: suffix, value: exponent
        current_exp = self.first_suffix_exp
        # Iterate over first suffices in each list
        for current_suffices in self.suffices:
            if not current_suffices: # No suffix to be used
                # Do not self suffix exp map
                self.exp_suffix_map[current_exp // 3] = ""
                current_exp += 3
                continue
            # Compute exponent -> suffix (only first suffix)
            self.exp_suffix_map[current_exp // 3] = current_suffices[0]
            # Compute suffix -> exponent
            for current_suffix in current_suffices:
                self.suffix_exp_map[current_suffix] = current_exp
            current_exp += 3
        # Compute min/max SI value
        self.exp_map_min = min(self.exp_suffix_map.keys())
        self.exp_map_max = max(self.exp_suffix_map.keys())

    def split_input(self, s):
        """
        Separate a string into a 3-tuple (number, suffix, unit).
        returns None if the string could not be parsed.

        The tuple will never contain None but empty strings if some
        element is not present. The number must be present for the string
        to be considered valid.

        Units always need to be a suffix. Instead
        Thousands separators or commata instead of points may be used
        (see normalize_interpunctation documentation).

        Thousands separators and suffix-as-decimal-separators may NOT
        be mixed. Whitespace is removed automatically.
        """
        # Remove thousands separator & ensure dot is used
        s = normalize_interpunctation(s).replace(" ", "")
        s, unit = self.split_unit(s) # Remove unit
        # Try to find SI suffix at the end or in the middle
        if s[-1] in self.all_suffixes:
            s, suffix = s[:-1], s[-1]
        else:  # Try to find unit anywhere
            isSuffixList = [ch in self.all_suffixes for ch in s]
            # Ensure only ONE unit occurs in the string
            suffixCount = sum(isSuffixList)
            if suffixCount > 1:
                raise ValueError("More than one SI suffix in the string")
            elif suffixCount == 0:
                suffix = ""
            else:  # suffixCount == 1 => correct
                # Suffix-as-decimal-separator --> no other decimal separator
                if "." in s:  # Commata already handled by normalize_interpunctation
                    return None
                suffixIndex = isSuffixList.index(True)
                # Suffix must NOT be first character
                if suffixIndex == 0:
                    raise ValueError("Suffix in '{0}' must not be the first char".format(s))
                suffix = s[suffixIndex]
                s = s.replace(suffix, ".")
        # Handle unit prefix (if any). Not allowable if no unit is present
        s = s.strip(self.strippable)
        # Final check: Is there any number left and is it valid?
        if not all((ch in _numeric_allowed for ch in s)):
            raise ValueError("Remainder of string is not purely numeric: {0}".format(s))
        return (s, suffix, unit)

    def split_unit(self, s):
        """
        Split a string into (remainder, unit).
        Only units in the units set are recognized
        unit may be "None" if no unit is recognized
        """
        # Fallback for strings which are too short
        if len(s) <= 1:
            return s, ""
        # Handle 2-character units (e.g. 'Hz'), then 1-character unity (e.g. 'V')
        if s[-2:] in self.units: # Will also handle unit-only 1-char strings
            s, unit = s[:-2], s[-2:]
        elif s[-1] in self.units:  # Handle 1-char units
            s, unit = s[:-1], s[-1]
        else: # No unit
            s, unit = s, ''
        # Remove unit prefix, if any (e.g. degrees symbol, delta symbol)
        s = s.rstrip(self.strippable)
        return s, unit

    def normalize(self, s, encoding="utf8"):
        """
        Converts an engineer's input of a wide variety of formats to a numeric
        value.

        Returns a pair (number, unit) or None if the conversion could not be performed.

        See splitSuffixSeparator() for further details on supported formats
        """
        # Make sure it's a decoded string
        if isinstance(s, bytes):
            s = s.decode(encoding)
        # Perform splitting
        res = self.split_input(s.strip())
        (num, suffix, unit) = res
        mul = (10 ** self.suffix_exp_map[suffix]) if suffix else 1
        return (float(num) * mul, unit)

    def safe_normalize(self, s, encoding="utf8"):
        """
        Same as normalize(), but returns None instead of raising
        on error.
        """
        try:
            return self.normalize(s, encoding)
        except ValueError:
            return None

    def format(self, v, unit=""):
        """
        Format v using SI suffices with optional units.
        Produces a string with 3 visible digits.
        """
        #Suffix map is indexed by one third of the decadic logarithm.
        exp = 0 if v == 0. else math.log(abs(v), 10.)
        suffixMapIdx = int(math.floor(exp / 3.))
        #Ensure we're in range
        if not self.exp_map_min < suffixMapIdx < self.exp_map_max:
            raise ValueError("Value out of range: {0}".format(v))
        #Pre-multiply the value
        v = v * (10.0 ** -(suffixMapIdx * 3))
        #Delegate the rest of the task to the helper
        return _formatWithSuffix(v, self.exp_suffix_map[suffixMapIdx] + unit)

    def auto_suffix_1d(self, arr):
        """
        Takes an array of arbitrary values and determines
        what is the best suffix (e.g. M, m, n, f) to represent
        as many values as possible with as few powers of 10 as possible.

        Returns a tuple (factor, suffix) where the factor is a floating-point
        value to multiply the array with to obtain value with "suffix" suffix.
        """
        # Compute logarithmic magnitudes of data
        arr_log = np.log10(np.abs(arr))
        arr_log[np.isinf(arr_log)] = 0  # log(0) == inf
        log_mean = arr_log.mean()
        # Generate score matrix
        suffix_idx = int(round(log_mean / 3))
        # Ensure we're in range
        suffix_idx = max(self.exp_map_min, suffix_idx)
        suffix_idx = min(self.exp_map_max, suffix_idx)
        # Pre-multiply the value
        multiplier = 10.0 ** -(suffix_idx * 3)
        return multiplier, self.exp_suffix_map[suffix_idx]

    def auto_format(self, fn, *args, **kwargs):
        """
        Auto-format a value by leveraging function annotations.
        For example this can be used to format using UliEngineering Physics quantities
        The function's return value is expected to a be annotated with a Quantity() value.
        """
        if not callable(fn):
            raise ValueError("fn must be callable")
        # Access innermost function inside possibly nested partials
        annotatedFN = fn
        while isinstance(annotatedFN, functools.partial):
            annotatedFN = annotatedFN.func
        try:
            qty = annotatedFN.__annotations__["return"]
        except KeyError:
            raise UnannotatedReturnValueError("Function {0} does not have an annotated return value")
        return self.format(fn(*args, **kwargs), qty.unit)

    def normalize_numeric_safe(self, arg):
        """
        Normalize each element of an iterable and retrieve only the numeric value
        (the unit is ignored). Works on iterables and string-likes.

        Use toolz.itertoolz.compact() on the result to remove all None values.

        Returns a list with None (on error) or the numeric value (no unit)
        """
        if isinstance(arg, int) or isinstance(arg, float):
            return arg
        # If it's stringlike, apply directly
        if isinstance(arg, str) or isinstance(arg, bytes):
            v = self.safe_normalize(arg)
            return None if v is None else v[0]
        # It's an iterable
        ret = []
        for elem in arg:
            v = self.safe_normalize(elem)
            ret.append(None if v is None else v[0])
        return ret

    def normalize_numeric(self, arg):
        """
        Normalize each element of an iterable and retrieve only the numeric value
        (the unit is ignored). Works on iterables and string-likes.

        Raises if any of the values can't be normalized
        """
        if isinstance(arg, int) or isinstance(arg, float):
            return arg
        # If it's stringlike, apply directly
        if isinstance(arg, str) or isinstance(arg, bytes):
            return self.normalize(arg)[0]
        # It's an iterable
        return [self.normalize(elem)[0] for elem in arg]

# Default instance, initialized on first use

def __init_engineer_io_instance():
    """
    Helper function to initialize the default instance
    the first time any of the global functions is used.
    """
    if EngineerIO.instance is None:
        EngineerIO.instance = EngineerIO()

__replace_comma_dot = lambda s: s.replace(",", ".")
"""
Map of a transform to apply to the string
during interpunctation normalization,
depending on (commaFound, dotFound, commaFoundFirst).
Must contain every possible variant
"""
_interpunct_transform_map = {
    # Found nothing or only point -> no modification required
    (False, False, False): functoolz.identity,
    (False, False, True): functoolz.identity,
    (False, True, False): functoolz.identity,
    (False, True, True): functoolz.identity,
    # Only comma -> replace and exit
    (True, False, False): __replace_comma_dot,
    (True, False, True): __replace_comma_dot,
    # Below this line: Both comma and dot found
    # Comma first => comma used as thousands separators
    (True, True, True): lambda s: s.replace(",", ""),
    # Dot first => dot used as thousands separator
    (True, True, False): lambda s: s.replace(".", "").replace(",", ".")
}

def normalize_interpunctation(s):
    """
    Normalize comma to point for float conversion.
    Correctly handles thousands separators.

    Note that cases like "1,234" are undecidable between
    "1234" and "1.234". They are treated as "1.234".

    Only points and commata are potentially modified.
    Other characters and digits are not handled.
    """
    foundCommaFirst = False
    commaIdx = s.find(",")
    pointIdx = s.find(".")
    foundComma = commaIdx is not None
    foundPoint = pointIdx is not None
    commaFirst = commaIdx < pointIdx if (foundComma and foundPoint) else None
    # Get the appropriate transform function from the map an run it on s
    return _interpunct_transform_map[(foundComma, foundPoint, commaFirst)](s)

def _formatWithSuffix(v, suffix=""):
    """
    Format a given value with a given suffix.
    This helper function formats the value to 3 visible digits.
    v must be pre-multiplied by the factor implied by the suffix
    """
    if v < 10.0:
        res = "{:.2f}".format(v)
    elif v < 100.0:
        res = "{:.1f}".format(v)
    else:  # Should only happen if v < 1000
        res = str(int(round(v)))
    #Avoid appending whitespace if there is no suffix
    return "{0} {1}".format(res, suffix) if suffix else res

def normalizeEngineerInput(s, encoding="utf8"):
    __init_engineer_io_instance()
    return EngineerIO.instance.normalize(s)

def formatValue(v, unit=""):
    __init_engineer_io_instance()
    return EngineerIO.instance.normalize(v, unit)

def safeNormalizeEngineerInput(v, unit=""):
    __init_engineer_io_instance()
    return EngineerIO.instance.safe_normalize(v, unit)

def normalize_numeric(v):
    __init_engineer_io_instance()
    return EngineerIO.instance.normalize_numeric(v)

def auto_format(v, *args, **kwargs):
    __init_engineer_io_instance()
    return EngineerIO.instance.auto_format(v, *args, **kwargs)
<|endoftext|>


<commit_before>import operator
import re

from functools import partial

from ..cache import Cache
from ..exceptions import PluginError, NoStreamsError
from ..options import Options

QUALITY_WEIGTHS_EXTRA = {
    "other": {
        "live": 1080,
    },
    "tv": {
        "hd": 1080,
        "sd": 576,
    },
    "quality": {
        "ehq": 720,
        "hq":  576,
        "sq":  360,
    },
}

FILTER_OPERATORS = {
    "<": operator.lt,
    "<=": operator.le,
    ">": operator.gt,
    ">=": operator.ge,
}

def stream_weight(stream):
    for group, weights in QUALITY_WEIGTHS_EXTRA.items():
        if stream in weights:
            return weights[stream], group

    match = re.match("^(\d+)([k]|[p])?([\+])?$", stream)

    if match:
        if match.group(2) == "k":
            bitrate = int(match.group(1))

            # FIXME: This is a crude attempt at making a bitrate's
            # weight end up similar to the weight of a resolution.
            # Someone who knows math, please fix.
            weight = bitrate / 2.8

            return weight, "bitrate"

        elif match.group(2) == "p":
            weight = int(match.group(1))

            if match.group(3) == "+":
                weight += 1

            return weight, "pixels"

    return 0, "none"

def iterate_streams(streams):
    for name, stream in streams.items():
        if isinstance(stream, list):
            for sub_stream in stream:
                yield (name, sub_stream)
        else:
            yield (name, stream)

def stream_type_priority(stream_types, stream):
    stream_type = type(stream[1]).shortname()

    try:
        prio = stream_types.index(stream_type)
    except ValueError:
        prio = 99

    return prio

def stream_sorting_filter(expr, stream_weight):
    match = re.match(r"(?P<op><=|>=|<|>)?(?P<value>[\w\+]+)", expr)

    if not match:
        raise PluginError("Invalid filter expression: {0}".format(expr))

    op, value = match.group("op", "value")
    op = FILTER_OPERATORS.get(op, operator.eq)
    filter_weight, filter_group = stream_weight(value)

    def func(quality):
        weight, group = stream_weight(quality)

        if group == filter_group:
            return not op(weight, filter_weight)

        return True

    return func

class Plugin(object):
    """A plugin can retrieve stream information from the URL specified.

    :param url: URL that the plugin will operate on
    """

    cache = None
    logger = None
    module = "unknown"
    options = Options()
    session = None

    @classmethod
    def bind(cls, session, module):
        cls.cache = Cache(filename="plugin-cache.json",
                          key_prefix=module)
        cls.logger = session.logger.new_module("plugin." + module)
        cls.module = module
        cls.session = session

    def __init__(self, url):
        self.url = url

    @classmethod
    def can_handle_url(cls, url):
        raise NotImplementedError

    @classmethod
    def set_option(cls, key, value):
        cls.options.set(key, value)

    @classmethod
    def get_option(cls, key):
        return cls.options.get(key)

    @classmethod
    def stream_weight(cls, stream):
        return stream_weight(stream)

    @classmethod
    def default_stream_types(cls, streams):
        stream_types = ["rtmp", "hls", "hds", "http"]

        for name, stream in iterate_streams(streams):
            stream_type = type(stream).shortname()

            if stream_type not in stream_types:
                stream_types.append(stream_type)

        return stream_types

    def get_streams(self, stream_types=None, sorting_excludes=None):
        """Attempts to extract available streams.

        Returns a :class:`dict` containing the streams, where the key is
        the name of the stream, most commonly the quality and the value
        is a :class:`Stream` object.

        The result can contain the synonyms **best** and **worst** which
        points to the streams which are likely to be of highest and
        lowest quality respectively.

        If multiple streams with the same name are found, the order of
        streams specified in *stream_types* will determine which stream
        gets to keep the name while the rest will be renamed to
        "<name>_<stream type>".

        The synonyms can be fine tuned with the *sorting_excludes*
        parameter. This can be either of these types:

            - A list of filter expressions in the format
              *[operator]<value>*. For example the filter ">480p" will
              exclude streams ranked higher than "480p" from the list
              used in the synonyms ranking. Valid operators are >, >=, <
              and <=. If no operator is specified then equality will be
              tested.

            - A function that is passed to filter() with a list of
              stream names as input.

        :param stream_types: A list of stream types to return.
        :param sorting_excludes: Specify which streams to exclude from
                                 the best/worst synonyms.

        .. versionchanged:: 1.4.2
           Added *priority* parameter.

        .. versionchanged:: 1.5.0
           Renamed *priority* to *stream_types* and changed behaviour
           slightly.

        .. versionchanged:: 1.5.0
           Added *sorting_excludes* parameter.

        .. versionchanged:: 1.6.0
           *sorting_excludes* can now be a list of filter expressions
           or a function that is passed to filter().

        """

        try:
            ostreams = self._get_streams()
        except NoStreamsError:
            return {}
        except (IOError, OSError, ValueError) as err:
            raise PluginError(err)

        if not ostreams:
            return {}

        streams = {}

        if stream_types is None:
            stream_types = self.default_stream_types(ostreams)

        # Add streams depending on stream type and priorities
        sorted_streams = sorted(iterate_streams(ostreams),
                                key=partial(stream_type_priority,
                                            stream_types))

        for name, stream in sorted_streams:
            stream_type = type(stream).shortname()

            if stream_type not in stream_types:
                continue

            if name in streams:
                name = "{0}_{1}".format(name, stream_type)

            # Validate stream name and discard the stream if it's bad.
            match = re.match("([A-z0-9_+]+)", name)
            if match:
                name = match.group(1)
            else:
                self.logger.debug("The stream '{0}' has been ignored "
                                  "since it is badly named.", name)
                continue

            # Force lowercase name and replace space with underscore.
            streams[name.lower()] = stream

        # Create the best/worst synonmys
        stream_weight_only = lambda s: (self.stream_weight(s)[0] or
                                        (len(streams) == 1 and 1))
        stream_names = filter(stream_weight_only, streams.keys())
        sorted_streams = sorted(stream_names, key=stream_weight_only)

        if isinstance(sorting_excludes, list):
            for expr in sorting_excludes:
                filter_func = stream_sorting_filter(expr, self.stream_weight)
                sorted_streams = list(filter(filter_func, sorted_streams))
        elif callable(sorting_excludes):
            sorted_streams = list(filter(sorting_excludes, sorted_streams))

        if len(sorted_streams) > 0:
            best = sorted_streams[-1]
            worst = sorted_streams[0]
            streams["best"] = streams[best]
            streams["worst"] = streams[worst]

        return streams

    def _get_streams(self):
        raise NotImplementedError

__all__ = ["Plugin"]
<commit_msg>api: Rename Plugin.get_streams to Plugin.streams.<commit_after>import operator
import re

from functools import partial

from ..cache import Cache
from ..exceptions import PluginError, NoStreamsError
from ..options import Options

QUALITY_WEIGTHS_EXTRA = {
    "other": {
        "live": 1080,
    },
    "tv": {
        "hd": 1080,
        "sd": 576,
    },
    "quality": {
        "ehq": 720,
        "hq":  576,
        "sq":  360,
    },
}

FILTER_OPERATORS = {
    "<": operator.lt,
    "<=": operator.le,
    ">": operator.gt,
    ">=": operator.ge,
}

def stream_weight(stream):
    for group, weights in QUALITY_WEIGTHS_EXTRA.items():
        if stream in weights:
            return weights[stream], group

    match = re.match("^(\d+)([k]|[p])?([\+])?$", stream)

    if match:
        if match.group(2) == "k":
            bitrate = int(match.group(1))

            # FIXME: This is a crude attempt at making a bitrate's
            # weight end up similar to the weight of a resolution.
            # Someone who knows math, please fix.
            weight = bitrate / 2.8

            return weight, "bitrate"

        elif match.group(2) == "p":
            weight = int(match.group(1))

            if match.group(3) == "+":
                weight += 1

            return weight, "pixels"

    return 0, "none"

def iterate_streams(streams):
    for name, stream in streams.items():
        if isinstance(stream, list):
            for sub_stream in stream:
                yield (name, sub_stream)
        else:
            yield (name, stream)

def stream_type_priority(stream_types, stream):
    stream_type = type(stream[1]).shortname()

    try:
        prio = stream_types.index(stream_type)
    except ValueError:
        prio = 99

    return prio

def stream_sorting_filter(expr, stream_weight):
    match = re.match(r"(?P<op><=|>=|<|>)?(?P<value>[\w\+]+)", expr)

    if not match:
        raise PluginError("Invalid filter expression: {0}".format(expr))

    op, value = match.group("op", "value")
    op = FILTER_OPERATORS.get(op, operator.eq)
    filter_weight, filter_group = stream_weight(value)

    def func(quality):
        weight, group = stream_weight(quality)

        if group == filter_group:
            return not op(weight, filter_weight)

        return True

    return func

class Plugin(object):
    """A plugin can retrieve stream information from the URL specified.

    :param url: URL that the plugin will operate on
    """

    cache = None
    logger = None
    module = "unknown"
    options = Options()
    session = None

    @classmethod
    def bind(cls, session, module):
        cls.cache = Cache(filename="plugin-cache.json",
                          key_prefix=module)
        cls.logger = session.logger.new_module("plugin." + module)
        cls.module = module
        cls.session = session

    def __init__(self, url):
        self.url = url

    @classmethod
    def can_handle_url(cls, url):
        raise NotImplementedError

    @classmethod
    def set_option(cls, key, value):
        cls.options.set(key, value)

    @classmethod
    def get_option(cls, key):
        return cls.options.get(key)

    @classmethod
    def stream_weight(cls, stream):
        return stream_weight(stream)

    @classmethod
    def default_stream_types(cls, streams):
        stream_types = ["rtmp", "hls", "hds", "http"]

        for name, stream in iterate_streams(streams):
            stream_type = type(stream).shortname()

            if stream_type not in stream_types:
                stream_types.append(stream_type)

        return stream_types

    def streams(self, stream_types=None, sorting_excludes=None):
        """Attempts to extract available streams.

        Returns a :class:`dict` containing the streams, where the key is
        the name of the stream, most commonly the quality and the value
        is a :class:`Stream` object.

        The result can contain the synonyms **best** and **worst** which
        points to the streams which are likely to be of highest and
        lowest quality respectively.

        If multiple streams with the same name are found, the order of
        streams specified in *stream_types* will determine which stream
        gets to keep the name while the rest will be renamed to
        "<name>_<stream type>".

        The synonyms can be fine tuned with the *sorting_excludes*
        parameter. This can be either of these types:

            - A list of filter expressions in the format
              *[operator]<value>*. For example the filter ">480p" will
              exclude streams ranked higher than "480p" from the list
              used in the synonyms ranking. Valid operators are >, >=, <
              and <=. If no operator is specified then equality will be
              tested.

            - A function that is passed to filter() with a list of
              stream names as input.

        :param stream_types: A list of stream types to return.
        :param sorting_excludes: Specify which streams to exclude from
                                 the best/worst synonyms.

        .. versionchanged:: 1.4.2
           Added *priority* parameter.

        .. versionchanged:: 1.5.0
           Renamed *priority* to *stream_types* and changed behaviour
           slightly.

        .. versionchanged:: 1.5.0
           Added *sorting_excludes* parameter.

        .. versionchanged:: 1.6.0
           *sorting_excludes* can now be a list of filter expressions
           or a function that is passed to filter().

        """

        try:
            ostreams = self._get_streams()
        except NoStreamsError:
            return {}
        except (IOError, OSError, ValueError) as err:
            raise PluginError(err)

        if not ostreams:
            return {}

        streams = {}

        if stream_types is None:
            stream_types = self.default_stream_types(ostreams)

        # Add streams depending on stream type and priorities
        sorted_streams = sorted(iterate_streams(ostreams),
                                key=partial(stream_type_priority,
                                            stream_types))

        for name, stream in sorted_streams:
            stream_type = type(stream).shortname()

            if stream_type not in stream_types:
                continue

            if name in streams:
                name = "{0}_{1}".format(name, stream_type)

            # Validate stream name and discard the stream if it's bad.
            match = re.match("([A-z0-9_+]+)", name)
            if match:
                name = match.group(1)
            else:
                self.logger.debug("The stream '{0}' has been ignored "
                                  "since it is badly named.", name)
                continue

            # Force lowercase name and replace space with underscore.
            streams[name.lower()] = stream

        # Create the best/worst synonmys
        stream_weight_only = lambda s: (self.stream_weight(s)[0] or
                                        (len(streams) == 1 and 1))
        stream_names = filter(stream_weight_only, streams.keys())
        sorted_streams = sorted(stream_names, key=stream_weight_only)

        if isinstance(sorting_excludes, list):
            for expr in sorting_excludes:
                filter_func = stream_sorting_filter(expr, self.stream_weight)
                sorted_streams = list(filter(filter_func, sorted_streams))
        elif callable(sorting_excludes):
            sorted_streams = list(filter(sorting_excludes, sorted_streams))

        if len(sorted_streams) > 0:
            best = sorted_streams[-1]
            worst = sorted_streams[0]
            streams["best"] = streams[best]
            streams["worst"] = streams[worst]

        return streams

    def get_streams(self, *args, **kwargs):
        """Deprecated since version 1.9.0.

        Has been renamed to :func:`Plugin.streams`, this is an alias
        for backwards compatibility.
        """

        return self.streams(*args, **kwargs)

    def _get_streams(self):
        raise NotImplementedError

__all__ = ["Plugin"]
<|endoftext|>


<commit_before>import re

from operator import methodcaller

from livestreamer.compat import urlparse
from livestreamer.plugin import Plugin
from livestreamer.plugin.api import http, validate
from livestreamer.stream import RTMPStream

_url_re = re.compile("http(s)?://(\w+\.)?ilive.to/")
_rtmp_re = re.compile("""
    \$.getJSON\("(?P<token_url>[^"]+)".*
    flashplayer:\s+"(?P<swf_url>[^"]+)".*
    streamer:\s+"(?P<rtmp_url>[^"]+)".*
    file:\s+"(?P<rtmp_playpath>[^"]+)\.flv"
""", re.VERBOSE | re.DOTALL)

_token_schema = validate.Schema(
    {
        "token": validate.text
    },
    validate.get("token")
)
_schema = validate.Schema(
    validate.transform(_rtmp_re.search),
    validate.any(
        None,
        validate.all(
            validate.transform(methodcaller("groupdict")),
            {
                "rtmp_playpath": validate.text,
                "rtmp_url": validate.all(
                    validate.transform(methodcaller("replace", "\\/", "/")),
                    validate.url(scheme="rtmp"),
                ),
                "swf_url": validate.url(scheme="http"),
                "token_url": validate.url(scheme="http")
            }
        )
    )
)

class ILive(Plugin):
    @classmethod
    def can_handle_url(self, url):
        return _url_re.match(url)

    def _get_streams(self):
        info = http.get(self.url, schema=_schema)
        if not info:
            return

        headers = {"Referer": self.url}
        res = http.get(info["token_url"], headers=headers)
        token = http.json(res, schema=_token_schema)

        parsed = urlparse(info["rtmp_url"])
        if parsed.query:
            app = "{0}?{1}".format(parsed.path[1:], parsed.query)
        else:
            app = parsed.path[1:]

        params = {
            "rtmp": info["rtmp_url"],
            "app": app,
            "pageUrl": self.url,
            "swfVfy": info["swf_url"],
            "playpath": info["rtmp_playpath"],
            "token": token,
            "live": True
        }

        stream = RTMPStream(self.session, params)
        return dict(live=stream)

__plugin__ = ILive
<commit_msg>plugins.ilive: Update for service changes.<commit_after>import re

from livestreamer.compat import urlparse
from livestreamer.plugin import Plugin
from livestreamer.plugin.api import StreamMapper, http, validate
from livestreamer.stream import HLSStream, RTMPStream

CHANNEL_URL = "http://www.mobileonline.tv/channel.php"

_url_re = re.compile("http(s)?://(\w+\.)?streamlive.to/view/(?P<channel>\d+)")
_link_re = re.compile("
<a href=(\S+) target=")
_schema = validate.Schema(
    validate.transform(_link_re.findall),
)

class ILive(Plugin):
    @classmethod
    def can_handle_url(self, url):
        return _url_re.match(url)

    def _create_hls_streams(self, url):
        try:
            streams = HLSStream.parse_variant_playlist(self.session, url)
            return streams.items()
        except IOError as err:
            self.logger.warning("Failed to extract HLS streams: {0}", err)

    def _create_rtmp_stream(self, url):
        parsed = urlparse(url)
        if parsed.query:
            app = "{0}?{1}".format(parsed.path[1:], parsed.query)
        else:
            app = parsed.path[1:]

        params = {
            "rtmp": url,
            "app": app,
            "pageUrl": self.url,
            "live": True
        }

        stream = RTMPStream(self.session, params)
        return "live", stream

    def _get_streams(self):
        channel = _url_re.match(self.url).group("channel")
        urls = http.get(CHANNEL_URL, params=dict(n=channel), schema=_schema)
        if not urls:
            return

        mapper = StreamMapper(cmp=lambda scheme, url: url.startswith(scheme))
        mapper.map("http", self._create_hls_streams)
        mapper.map("rtmp", self._create_rtmp_stream)

        return mapper(urls)

__plugin__ = ILive
<|endoftext|>