<commit_before>#!/usr/bin/python3 -u

import sys
import os
from dbus.mainloop.glib import DBusGMainLoop
import dbus
import dbus.service
from gi.repository import GLib
import pickle
import logging
from collections import defaultdict
from functools import total_ordering
from itertools import repeat
import heapq
from argparse import ArgumentParser
import json
from csv import reader as csvreader
from csv import excel_tab

InterfaceBusItem = 'com.victronenergy.BusItem'
InterfaceProperties = 'org.freedesktop.DBus.Properties'

class SystemBus(dbus.bus.BusConnection):
    def __new__(cls):
        return dbus.bus.BusConnection.__new__(cls, dbus.bus.BusConnection.TYPE_SYSTEM)

class SessionBus(dbus.bus.BusConnection):
    def __new__(cls):
        return dbus.bus.BusConnection.__new__(cls, dbus.bus.BusConnection.TYPE_SESSION)

def wrap_dbus_value(v):
	""" Stripped down version, since everything is already dbus wrapped, we
	    just need to handle empty lists. """
	if isinstance(v, list) and len(v) == 0:
		return dbus.Array([], signature=dbus.Signature('u'), variant_level=1)
	return v

# Used for CSV to dbus wrapping
def demarshall(typ, v):
	de = {
		"INT32": dbus.types.Int32,
		"UINT32": dbus.types.UInt32,
		"UINT16": dbus.types.UInt16,
		"BYTE": lambda x: dbus.types.Byte(int(x)),
		"DOUBLE": dbus.types.Double,
		"STRING": dbus.types.String,
	}

	if typ.startswith("ARRAY"):
		subtype = typ[6:-1]
		v = json.loads(v)
		return dbus.types.Array([demarshall(subtype, x) for x in v])

	if typ in de:
		return de[typ](v)

	return None

class Unpickler(pickle.Unpickler):
	# Support for older pickles
	_replacements = {
		('dbus', 'UTF8String'): ('dbus', 'String')
	}
	def find_class(self, module, name):
		module, name = self._replacements.get((module, name), (module, name))
		return super().find_class(module, name)

class DbusRootObject(dbus.service.Object):
	_objectPath = '/'

	def __init__(self, busName, values):
		super(DbusRootObject, self).__init__(busName, '/')
		self.values = values

	@dbus.service.method(InterfaceBusItem, out_signature = 'v')
	def GetValue(self):
		values = { k[1:]: wrap_dbus_value(v._properties['Value']) \
			for k, v in list(self.values.items()) \
			if hasattr(v, '_properties') and 'Value' in v._properties }
		return dbus.Dictionary(values, signature=dbus.Signature('sv'),
			variant_level=1)

	def setItems(self, items):
		# Pass on to DbusPathObject storage
		for path, changes in items.items():
			try:
				self.values[path]._setProperties(changes)
			except KeyError:
				pass

		self.notify(items)

	def notify(self, items):
		# Send ItemsChanged. Make sure it has the right signature
		# since python-dbus will guess wrong.
		self.ItemsChanged({ k: dbus.types.Dictionary(v, signature='sv')
			for k, v in items.items() })

	@dbus.service.signal(InterfaceBusItem, signature='a{sa{sv}}')
	def ItemsChanged(self, items):
		logging.debug('signal ItemsChanged')

class LegacyDbusPathObject(dbus.service.Object):
	## Constructor of MyDbusObject
	#
	# Creates the dbus-object under the given bus-name (dbus-service-name).
	# @param busName Return value from dbus.service.BusName, see run()).
	# @param objectPath The dbus-object-path.
	def __init__(self, busName, objectPath, properties):
		super(LegacyDbusPathObject, self).__init__(busName, objectPath)
		self._objectPath = objectPath
		self._properties = properties

	@dbus.service.method(InterfaceBusItem, out_signature = 'v')
	def GetValue(self):
		logging.debug("GetValue %s" % self._objectPath)
		value = dbus.Array([], signature=dbus.Signature('i'), variant_level=1)
		if 'Value' in self._properties:
			if self._properties['Value'] != dbus.Array([]):
				value = self._properties['Value']
		return value

	## Dbus method GetText
	# Returns the value as string of the dbus-object-path.
	# @return text A text-value or '' (error)
	@dbus.service.method(InterfaceBusItem, out_signature = 's')
	def GetText(self):
		logging.debug("GetText %s" % self._objectPath)
		if 'Text' in self._properties:
			return self._properties['Text']
		return '---'

	## Dbus method SetValue
	# Sets the value.
	# @param value The new value.
	# @return completion-code When successful a 0 is return, and when not a -1 is returned.
	@dbus.service.method(InterfaceBusItem, in_signature = 'v', out_signature = 'i')
	def SetValue(self, value):
		logging.debug("SetValue %s" % self._objectPath)
		if 'Value' in self._properties:
			self._properties['Value'] = value
			self._properties['Text'] = str(value)
			self.notify(self._properties)
			return 0
		return -1

	def _setProperties(self, properties):
		# Make sure it is wrapped. If no Value, substitute invalid.
		properties['Value'] = wrap_dbus_value(properties.get('Value', []))
		self._properties = properties

	def setProperties(self, properties):
		self._setProperties(properties)
		self.notify(properties)

	def notify(self, properties):
		self.PropertiesChanged(properties)

	@dbus.service.signal(InterfaceBusItem, signature = 'a{sv}')
	def PropertiesChanged(self, properties):
		logging.debug('signal PropertiesChanged %s %s' % (self._object_path, properties))

class DbusPathObject(LegacyDbusPathObject):
	""" Newer style that uses ItemsChanged. """
	def __init__(self, root, busName, objectPath, properties):
		super(DbusPathObject, self).__init__(busName, objectPath, properties)
		self._root = root

	def notify(self, properties):
		# Send an ItemsChanged on root instead
		self._root.notify({ self._objectPath: properties })

@total_ordering
class PropertiesChangedData(object):
	def __init__(self, time, dbusObjectPath, changes):
		self._time = time
		self._dbusObjectPath = dbusObjectPath
		self._changes = changes

	def __eq__(self, other):
		return self._time == other._time

	def __lt__(self, other):
		return self._time < other._time

class PickleIterator(object):
	def __init__(self, unpickler, service, values):
		self.unpickler = unpickler
		self.service = service
		self.values = values

	def __iter__(self):
		return self

	def __next__(self):
		try:
			return self.unpickler.load()
		except EOFError:
			raise StopIteration

class CsvIterator(object):
	def __init__(self, csvreader, service, values):
		self.csvreader = csvreader
		self.service = service # com.victronenergy.battery.whatever
		self.values = values # /Dc/0/Voltage: {'Value': 53.2, 'Text': '53.2V'}

	def __iter__(self):
		return self

	def __next__(self):
		try:
			t, path, typ, value, text = next(self.csvreader)[:5]
			return PropertiesChangedData(int(t), path, {
				'Value': demarshall(typ, value),
				'Text': text
			})
		except (EOFError, ValueError):
			raise StopIteration

class EventStream(object):
	def __init__(self, *args):
		self.args = args
		self.reset()

	def reset(self):
		self.streams = [x for x in (open_recording(fn) for fn in self.args) if x is not None]
		self.events = sort_streams(*self.streams)
		self.current = None

	def __iter__(self):
		return self

	def __next__(self):
		self.current = next(self.events)
		return self.current

class Timer(object):
	def __init__(self, services, events, nozip):
		self.services = services
		self.events = events
		self.nozip = nozip
		self.tickcount = 0
		self.evt = None

	def current_event(self):
		if self.events.current is None:
			return next(self.events)
		return self.events.current

	def send_unzipped(self):
		evt = self.current_event()
		while evt[0]._time <= self.tickcount:
			servicename = evt[1]
			servicepath = evt[0]._dbusObjectPath
			service = self.services[servicename]
			if servicepath in service:
				changes = evt[0]._changes
				logging.debug('Replaying %s %s %s' % (servicename, servicepath, changes))
				service[servicepath].setProperties(changes)
			evt = next(self.events)

	def send_zipped(self):
		di = defaultdict(dict)
		evt = self.current_event()
		while evt[0]._time <= self.tickcount:
			servicename = evt[1]
			servicepath = evt[0]._dbusObjectPath
			service = self.services[servicename]
			if servicepath in service:
				di[servicename][servicepath] = evt[0]._changes
			evt = next(self.events)

		for n, i in di.items():
			logging.debug('Replaying %s / %s' % (n, i))
			service = self.services[n]
			service['/'].setItems(i)

	def __call__(self):
		evt = self.current_event()
		try:
			if self.nozip:
				self.send_unzipped()
			else:
				self.send_zipped()
		except StopIteration:
			# Reset to start of simulation state
			if self.nozip:
				for it in self.events.streams:
					for path, props in it.values.items():
						if self.services[it.service][path].GetValue() != props['Value']:
							self.services[it.service][path].setProperties(props)
			else:
				for it in self.events.streams:
					di = {}
					for path, props in it.values.items():
						if self.services[it.service][path].GetValue() != props['Value']:
							di[path] = props
					if di:
						self.services[it.service]['/'].setItems(di)

			# Restart event stream (changes)
			self.events.reset()
			self.tickcount = 0
			evt = next(self.events)
		else:
			self.tickcount += 1

		return True

def open_recording(fn):
	if fn.endswith('.csv'):
		fp = open(fn, 'r', encoding='UTF-8')
		reader = csvreader(fp, dialect=excel_tab, quotechar="'")
		try:
			service = next(reader)[0]
		except IndexError:
			fp.close()
			return None

		try:
			values = {}
			path, typ, value, text = next(reader)[:4]
			while path:
				values[path] = {
					'Value': demarshall(typ, value),
					'Text': text
				}
				path, typ, value, text = next(reader)[:4]
		except EOFError:
			fp.close()
			return None
		except ValueError:
			pass # Not enough values to unpack

		return CsvIterator(reader, service, values)
	else:
		fp = open(fn, 'rb')
		unpickler = Unpickler(fp, encoding='UTF-8')
		try:
			service = unpickler.load()
			values = unpickler.load()
		except EOFError:
			fp.close()
			return None

		return PickleIterator(unpickler, service, values)

def sort_streams(*args):
	""" args is a list of iterables, The smallest item is yielded each time.
	"""
	streams = [zip(c, repeat(c.service)) for c in args]
	return heapq.merge(*streams)

def main():
	parser = ArgumentParser(description=sys.argv[0])
	parser.add_argument('--speed',
		help='Speedup factor', type=int, default=1)
	parser.add_argument('--nozip',
		action="store_true",
		help='Do not zip multiple PropertiesChanged into a single ItemsChanged',
		default=False)
	parser.add_argument('datafiles', nargs='+', help='Path to data file')

	args = parser.parse_args()

	DBusGMainLoop(set_as_default=True)

	services = defaultdict(dict)
	if args.nozip:
		ObjectFactory = lambda r, *args: LegacyDbusPathObject(*args)
	else:
		ObjectFactory = DbusPathObject

	logging.info("Opening recordings")
	events = EventStream(*args.datafiles)

	# Register dbus services/paths
	for it in events.streams:
		# Every stream gets its own dbus connection.
		logging.info("Initialising service {}".format(it.service))
		bus = SessionBus() if 'DBUS_SESSION_BUS_ADDRESS' in os.environ else SystemBus()
		busName = dbus.service.BusName(it.service, bus)
		service = services[it.service]

		# Root object
		service['/'] = root = DbusRootObject(busName, service)

		for path, props in it.values.items():
			logging.debug("path %s properties %s" % (path, props))
			service[path] = ObjectFactory(root, busName, path, props)

	logging.info("Stage set, starting simulation")
	GLib.timeout_add(int(1000.0/args.speed), Timer(services, events, args.nozip))
	GLib.MainLoop().run()

if __name__ == "__main__":
	logging.basicConfig(level=logging.INFO)
	main()
<commit_msg>play: support gzipped demo files<commit_after>#!/usr/bin/python3 -u

import sys
import os
from dbus.mainloop.glib import DBusGMainLoop
import dbus
import dbus.service
from gi.repository import GLib
import pickle
import logging
from collections import defaultdict
from functools import total_ordering
from itertools import repeat
import heapq
from argparse import ArgumentParser
import json
from csv import reader as csvreader
from csv import excel_tab
import gzip

InterfaceBusItem = 'com.victronenergy.BusItem'
InterfaceProperties = 'org.freedesktop.DBus.Properties'

class SystemBus(dbus.bus.BusConnection):
    def __new__(cls):
        return dbus.bus.BusConnection.__new__(cls, dbus.bus.BusConnection.TYPE_SYSTEM)

class SessionBus(dbus.bus.BusConnection):
    def __new__(cls):
        return dbus.bus.BusConnection.__new__(cls, dbus.bus.BusConnection.TYPE_SESSION)

def wrap_dbus_value(v):
	""" Stripped down version, since everything is already dbus wrapped, we
	    just need to handle empty lists. """
	if isinstance(v, list) and len(v) == 0:
		return dbus.Array([], signature=dbus.Signature('u'), variant_level=1)
	return v

# Used for CSV to dbus wrapping
def demarshall(typ, v):
	de = {
		"INT32": dbus.types.Int32,
		"UINT32": dbus.types.UInt32,
		"UINT16": dbus.types.UInt16,
		"BYTE": lambda x: dbus.types.Byte(int(x)),
		"DOUBLE": dbus.types.Double,
		"STRING": dbus.types.String,
	}

	if typ.startswith("ARRAY"):
		subtype = typ[6:-1]
		v = json.loads(v)
		return dbus.types.Array([demarshall(subtype, x) for x in v])

	if typ in de:
		return de[typ](v)

	return None

class Unpickler(pickle.Unpickler):
	# Support for older pickles
	_replacements = {
		('dbus', 'UTF8String'): ('dbus', 'String')
	}
	def find_class(self, module, name):
		module, name = self._replacements.get((module, name), (module, name))
		return super().find_class(module, name)

class DbusRootObject(dbus.service.Object):
	_objectPath = '/'

	def __init__(self, busName, values):
		super(DbusRootObject, self).__init__(busName, '/')
		self.values = values

	@dbus.service.method(InterfaceBusItem, out_signature = 'v')
	def GetValue(self):
		values = { k[1:]: wrap_dbus_value(v._properties['Value']) \
			for k, v in list(self.values.items()) \
			if hasattr(v, '_properties') and 'Value' in v._properties }
		return dbus.Dictionary(values, signature=dbus.Signature('sv'),
			variant_level=1)

	def setItems(self, items):
		# Pass on to DbusPathObject storage
		for path, changes in items.items():
			try:
				self.values[path]._setProperties(changes)
			except KeyError:
				pass

		self.notify(items)

	def notify(self, items):
		# Send ItemsChanged. Make sure it has the right signature
		# since python-dbus will guess wrong.
		self.ItemsChanged({ k: dbus.types.Dictionary(v, signature='sv')
			for k, v in items.items() })

	@dbus.service.signal(InterfaceBusItem, signature='a{sa{sv}}')
	def ItemsChanged(self, items):
		logging.debug('signal ItemsChanged')

class LegacyDbusPathObject(dbus.service.Object):
	## Constructor of MyDbusObject
	#
	# Creates the dbus-object under the given bus-name (dbus-service-name).
	# @param busName Return value from dbus.service.BusName, see run()).
	# @param objectPath The dbus-object-path.
	def __init__(self, busName, objectPath, properties):
		super(LegacyDbusPathObject, self).__init__(busName, objectPath)
		self._objectPath = objectPath
		self._properties = properties

	@dbus.service.method(InterfaceBusItem, out_signature = 'v')
	def GetValue(self):
		logging.debug("GetValue %s" % self._objectPath)
		value = dbus.Array([], signature=dbus.Signature('i'), variant_level=1)
		if 'Value' in self._properties:
			if self._properties['Value'] != dbus.Array([]):
				value = self._properties['Value']
		return value

	## Dbus method GetText
	# Returns the value as string of the dbus-object-path.
	# @return text A text-value or '' (error)
	@dbus.service.method(InterfaceBusItem, out_signature = 's')
	def GetText(self):
		logging.debug("GetText %s" % self._objectPath)
		if 'Text' in self._properties:
			return self._properties['Text']
		return '---'

	## Dbus method SetValue
	# Sets the value.
	# @param value The new value.
	# @return completion-code When successful a 0 is return, and when not a -1 is returned.
	@dbus.service.method(InterfaceBusItem, in_signature = 'v', out_signature = 'i')
	def SetValue(self, value):
		logging.debug("SetValue %s" % self._objectPath)
		if 'Value' in self._properties:
			self._properties['Value'] = value
			self._properties['Text'] = str(value)
			self.notify(self._properties)
			return 0
		return -1

	def _setProperties(self, properties):
		# Make sure it is wrapped. If no Value, substitute invalid.
		properties['Value'] = wrap_dbus_value(properties.get('Value', []))
		self._properties = properties

	def setProperties(self, properties):
		self._setProperties(properties)
		self.notify(properties)

	def notify(self, properties):
		self.PropertiesChanged(properties)

	@dbus.service.signal(InterfaceBusItem, signature = 'a{sv}')
	def PropertiesChanged(self, properties):
		logging.debug('signal PropertiesChanged %s %s' % (self._object_path, properties))

class DbusPathObject(LegacyDbusPathObject):
	""" Newer style that uses ItemsChanged. """
	def __init__(self, root, busName, objectPath, properties):
		super(DbusPathObject, self).__init__(busName, objectPath, properties)
		self._root = root

	def notify(self, properties):
		# Send an ItemsChanged on root instead
		self._root.notify({ self._objectPath: properties })

@total_ordering
class PropertiesChangedData(object):
	def __init__(self, time, dbusObjectPath, changes):
		self._time = time
		self._dbusObjectPath = dbusObjectPath
		self._changes = changes

	def __eq__(self, other):
		return self._time == other._time

	def __lt__(self, other):
		return self._time < other._time

class PickleIterator(object):
	def __init__(self, unpickler, service, values):
		self.unpickler = unpickler
		self.service = service
		self.values = values

	def __iter__(self):
		return self

	def __next__(self):
		try:
			return self.unpickler.load()
		except EOFError:
			raise StopIteration

class CsvIterator(object):
	def __init__(self, csvreader, service, values):
		self.csvreader = csvreader
		self.service = service # com.victronenergy.battery.whatever
		self.values = values # /Dc/0/Voltage: {'Value': 53.2, 'Text': '53.2V'}

	def __iter__(self):
		return self

	def __next__(self):
		try:
			t, path, typ, value, text = next(self.csvreader)[:5]
			return PropertiesChangedData(int(t), path, {
				'Value': demarshall(typ, value),
				'Text': text
			})
		except (EOFError, ValueError):
			raise StopIteration

class EventStream(object):
	def __init__(self, *args):
		self.args = args
		self.reset()

	def reset(self):
		self.streams = [x for x in (open_recording(fn) for fn in self.args) if x is not None]
		self.events = sort_streams(*self.streams)
		self.current = None

	def __iter__(self):
		return self

	def __next__(self):
		self.current = next(self.events)
		return self.current

class Timer(object):
	def __init__(self, services, events, nozip):
		self.services = services
		self.events = events
		self.nozip = nozip
		self.tickcount = 0
		self.evt = None

	def current_event(self):
		if self.events.current is None:
			return next(self.events)
		return self.events.current

	def send_unzipped(self):
		evt = self.current_event()
		while evt[0]._time <= self.tickcount:
			servicename = evt[1]
			servicepath = evt[0]._dbusObjectPath
			service = self.services[servicename]
			if servicepath in service:
				changes = evt[0]._changes
				logging.debug('Replaying %s %s %s' % (servicename, servicepath, changes))
				service[servicepath].setProperties(changes)
			evt = next(self.events)

	def send_zipped(self):
		di = defaultdict(dict)
		evt = self.current_event()
		while evt[0]._time <= self.tickcount:
			servicename = evt[1]
			servicepath = evt[0]._dbusObjectPath
			service = self.services[servicename]
			if servicepath in service:
				di[servicename][servicepath] = evt[0]._changes
			evt = next(self.events)

		for n, i in di.items():
			logging.debug('Replaying %s / %s' % (n, i))
			service = self.services[n]
			service['/'].setItems(i)

	def __call__(self):
		evt = self.current_event()
		try:
			if self.nozip:
				self.send_unzipped()
			else:
				self.send_zipped()
		except StopIteration:
			# Reset to start of simulation state
			if self.nozip:
				for it in self.events.streams:
					for path, props in it.values.items():
						if self.services[it.service][path].GetValue() != props['Value']:
							self.services[it.service][path].setProperties(props)
			else:
				for it in self.events.streams:
					di = {}
					for path, props in it.values.items():
						if self.services[it.service][path].GetValue() != props['Value']:
							di[path] = props
					if di:
						self.services[it.service]['/'].setItems(di)

			# Restart event stream (changes)
			self.events.reset()
			self.tickcount = 0
			evt = next(self.events)
		else:
			self.tickcount += 1

		return True

def open_recording(fn):
	_fn = fn
	if _fn.endswith('.gz'):
		_open = gzip.open
		_fn = fn[:-3]
	else:
		_open = open

	if _fn.endswith('.csv'):
		fp = _open(fn, 'rt', encoding='UTF-8')
		reader = csvreader(fp, dialect=excel_tab, quotechar="'")
		try:
			service = next(reader)[0]
		except IndexError:
			fp.close()
			return None

		try:
			values = {}
			path, typ, value, text = next(reader)[:4]
			while path:
				values[path] = {
					'Value': demarshall(typ, value),
					'Text': text
				}
				path, typ, value, text = next(reader)[:4]
		except EOFError:
			fp.close()
			return None
		except ValueError:
			pass # Not enough values to unpack

		return CsvIterator(reader, service, values)
	else:
		fp = _open(fn, 'rb')
		unpickler = Unpickler(fp, encoding='UTF-8')
		try:
			service = unpickler.load()
			values = unpickler.load()
		except EOFError:
			fp.close()
			return None

		return PickleIterator(unpickler, service, values)

def sort_streams(*args):
	""" args is a list of iterables, The smallest item is yielded each time.
	"""
	streams = [zip(c, repeat(c.service)) for c in args]
	return heapq.merge(*streams)

def main():
	parser = ArgumentParser(description=sys.argv[0])
	parser.add_argument('--speed',
		help='Speedup factor', type=int, default=1)
	parser.add_argument('--nozip',
		action="store_true",
		help='Do not zip multiple PropertiesChanged into a single ItemsChanged',
		default=False)
	parser.add_argument('datafiles', nargs='+', help='Path to data file')

	args = parser.parse_args()

	DBusGMainLoop(set_as_default=True)

	services = defaultdict(dict)
	if args.nozip:
		ObjectFactory = lambda r, *args: LegacyDbusPathObject(*args)
	else:
		ObjectFactory = DbusPathObject

	logging.info("Opening recordings")
	events = EventStream(*args.datafiles)

	# Register dbus services/paths
	for it in events.streams:
		# Every stream gets its own dbus connection.
		logging.info("Initialising service {}".format(it.service))
		bus = SessionBus() if 'DBUS_SESSION_BUS_ADDRESS' in os.environ else SystemBus()
		busName = dbus.service.BusName(it.service, bus)
		service = services[it.service]

		# Root object
		service['/'] = root = DbusRootObject(busName, service)

		for path, props in it.values.items():
			logging.debug("path %s properties %s" % (path, props))
			service[path] = ObjectFactory(root, busName, path, props)

	logging.info("Stage set, starting simulation")
	GLib.timeout_add(int(1000.0/args.speed), Timer(services, events, args.nozip))
	GLib.MainLoop().run()

if __name__ == "__main__":
	logging.basicConfig(level=logging.INFO)
	main()
<|endoftext|>


<commit_before>#!/bin/python
import sys
import os
import copy
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

def save(path, ext='png', close=True, verbose=True):
    """Save a figure from pyplot.

    Parameters
    ----------
    path : string
        The path (and filename, without the extension) to save the
        figure to.

    ext : string (default='png')
        The file extension. This must be supported by the active
        matplotlib backend (see matplotlib.backends module).  Most
        backends support 'png', 'pdf', 'ps', 'eps', and 'svg'.

    close : boolean (default=True)
        Whether to close the figure after saving.  If you want to save
        the figure multiple times (e.g., to multiple formats), you
        should NOT close it in between saves or you will have to
        re-plot it.

    verbose : boolean (default=True)
        Whether to print information about when and where the image
        has been saved.

    """

    # Extract the directory and filename from the given path
    directory = os.path.split(path)[0]
    filename = "%s.%s" % (os.path.split(path)[1], ext)
    if directory == '':
        directory = '.'

    # If the directory does not exist, create it
    if not os.path.exists(directory):
        os.makedirs(directory)

    # The final path to save to
    savepath = os.path.join(directory, filename)

    if verbose:
        print("Saving figure to '%s'..." % savepath),

    # Actually save the figure
    plt.savefig(savepath)

    # Close it
    if close:
        plt.close()

    if verbose:
        print("Done")

def read_dataframe_from_excel(
        filename,
        worksheet,
        **kwargs):
    """Reads a pandas dataframe from an Excel file
    filename: the excel file
    worksheet: worksheet name within the excel file
    The remaining named arguments are passed to pandas.read_excel
    """
    results = pd.read_excel(filename, worksheet, **kwargs)
    return results

def calculate_speedup_and_efficiency(
        rdf,
        compute_element_col_index,
        time_col_index):
    """Calculates the speedup and efficiency and
    adds them as new columns to a dataframe

    rdf: the results dataframe
    compute_element_col_index: compute elements column name
    time_col_index: column name containing the computation times

    Returns: a tuple containing the speedup and efficiency series
    """
    # get the t1 reference value
    t1 = float(rdf[rdf[compute_element_col_index] == 1][time_col_index])
    # speedup n = t1 / tn
    speedup = t1 / rdf[time_col_index]
    # efficiency n = t1 / (tn * n)
    efficiency = t1 / (rdf[time_col_index] * rdf[compute_element_col_index])

    return (speedup, efficiency)

def plot_walltime(
        series,
        colours,
        series_names,
        compute_elements,
        ymax,
        plot_size=None,
        group_width=0.8,
        xlabel='Processes',
        ylabel='Minutes',
        title='Walltime',
        file_name='walltime',
        file_extension='png',
        y_log_scale=False,
        show=False):
    """creates a bar plot as a new pyplot figure"""

    # the x locations for the groups.
    x_ind = np.arange(len(compute_elements))

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # create the individual bars
    bars_per_group = len(series)
    bar_width = group_width / bars_per_group
    bar_left = x_ind - (group_width / 2)  # centre the group on the ticks
    plot_left = x_ind[0] - group_width
    plot_right = x_ind[-1] + group_width

    for bar_values, bar_colour in zip(series, colours):
        ax.bar(
            bar_left,
            bar_values,
            width=bar_width,
            color=bar_colour,
            log=y_log_scale)
        bar_left += bar_width

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(0, ymax)
    ax.set_xticks(x_ind)
    ax.set_xticklabels(compute_elements)
    ax.legend(series_names)
    # ax.legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def plot_speedup(
        series,
        colours,
        series_names,
        compute_elements,
        xmax,
        ymax,
        line_width=1,
        plot_size=None,
        xlabel='Processes',
        ylabel='Speedup',
        title='Speedup',
        file_name='speedup',
        file_extension='png',
        show=False):
    """creates a speedup plot"""

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # define the sizes and locations of things
    x_ticks = copy.deepcopy(compute_elements)
    x_ticks.append(xmax)
    plot_left = x_ticks[0]
    plot_right = x_ticks[-1]
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(plot_left, ymax)

    # create the ideal speedup reference line
    ax.plot(
        x_ticks,
        x_ticks,
        color='red',
        linewidth=line_width,
        label='_nolegend_')

    # plot each series
    for values, colour, name in zip(series, colours, series_names):
        ax.plot(
            compute_elements,
            values,
            label=name,
            color=colour,
            linewidth=line_width,
            marker='o')

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(compute_elements)
    ax.set_xticklabels(compute_elements)
    ax.legend()
    ax.legend().get_frame().set_facecolor('white')
    # can also use get_frame().set_xy(x,y) to give a custom location

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def plot_efficiency(
        series,
        colours,
        series_names,
        compute_elements,
        xmax,
        ymax=1.2,
        line_width=1,
        plot_size=None,
        xlabel='Processes',
        ylabel='Efficiency',
        title='Efficiency',
        file_name='efficiency',
        file_extension='png',
        show=False):
    """creates an efficiency plot"""

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # define the sizes and locations of things
    x_ticks = copy.deepcopy(compute_elements)
    x_ticks.append(xmax)
    plot_left = x_ticks[0]
    plot_right = x_ticks[-1]
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(0, ymax)

    # create a horizontal ideal efficiency reference line
    ax.hlines(
        1,
        plot_left,
        plot_right,
        colors='red',
        linestyles='solid',
        linewidth=line_width,
        label='_nolegend_')

    # plot each series
    for values, colour, name in zip(series, colours, series_names):
        ax.plot(
            compute_elements,
            values,
            label=name,
            color=colour,
            linewidth=line_width,
            marker='o')

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(compute_elements)
    ax.set_xticklabels(compute_elements)
    ax.legend(shadow=True)
    ax.legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

if __name__ == '__main__':

    # create plots in a 4:3 aspect ratio
    # matplotlib works in inches
    plot_width = 10
    plot_size = (plot_width, plot_width * 3 / 4)

    # show_instead_of_save = False
    show_instead_of_save = True

    # CSIRO colours
    colours = [
        '#00a9ce',  # midday blue
        '#00616c',  # midnight blue
        '#78be20',  # light forest
        '#DF1995',  # fuschia
        '#E87722',  # orange
        '#E4002B',  # vermillion
        '#FFB81C',  # gold
        '#6D2077',  # plum
        '#1E22AA',  # blueberry
    ]

    results = read_dataframe_from_excel(
        'results.xls',
        worksheet='results',
        usecols=[
            'group',
            'compute_elements',
            'walltime'])

    # I don't know an equivalent to R's completecases, but this works for now
    results = results[results['group'].notnull() & (results.walltime > 0)]

    # create new dataframes for each group of results
    groups = results.group.unique()
    group_dataframes = {}
    for group in groups:
        df = results[(results.group == group)]
        if len(df):
            group_dataframes[group] = df.sort('compute_elements', ascending=True)

    # calculate speedup and efficiency based on the total walltime, and on
    # the average time per iteration (which excludes a lot of the serial
    # code, and gives a closer look at just the model calculations)
    for data in group_dataframes.values():
        speedup, efficiency = calculate_speedup_and_efficiency(
            data,
            'compute_elements',
            'walltime')
        data['speedup'] = speedup
        data['efficiency'] = efficiency

    # extract the compute element names for grouping and labelling the charts
    compute_elements = np.sort(results.compute_elements.unique()).tolist()

    # build the collections of values for the charts
    walltimes = []
    efficiencies = []
    speedups = []
    series_names = []
    max_walltime = 0
    for name, data in group_dataframes.items():
        walltimes.append(data.walltime)
        efficiencies.append(data.efficiency)
        speedups.append(data.speedup)
        series_names.append(name)
        local_max_walltime = max(data.walltime)
        if local_max_walltime > max_walltime:
            max_walltime = local_max_walltime

    # finally make the plots

    plot_walltime(
        walltimes,
        colours,
        series_names,
        compute_elements,
        ymax=max_walltime + 50,
        group_width=0.83,
        plot_size=plot_size,
        title='ODPS Walltime',
        xlabel='Threads',
        ylabel='Seconds',
        show=show_instead_of_save,
        file_name='odps-walltime',
        file_extension='png')

    plot_efficiency(
        efficiencies,
        colours,
        series_names,
        compute_elements,
        line_width=1.5,
        xmax=140,
        ymax=1.3,
        plot_size=plot_size,
        title='ODPS - Strong Scaling Efficiency',
        show=show_instead_of_save,
        file_name='odps-efficiency',
        file_extension='png')

    # plot_speedup(
        # speedups,
        # colours,
        # series_names,
        # compute_elements,
        # line_width=1.5,
        # xmax=140,
        # ymax=140,
        # plot_size=plot_size,
        # title='Gridded Calibration - Strong Scaling Speedup',
        # show=show_instead_of_save,
        # file_name='speedup',
        # file_extension='png')

    # plot_efficiency(
        # efficiencies_it,
        # colours,
        # series_names,
        # compute_elements,
        # line_width=1.5,
        # xmax=140,
        # ymax=1.5,
        # plot_size=plot_size,
        # title='Gridded Calibration - Strong Scaling Efficiency - Model Calculations',
        # show=show_instead_of_save,
        # file_name='efficiency_iterations',
        # file_extension='png')

    # plot_speedup(
        # speedups_it,
        # colours,
        # series_names,
        # compute_elements,
        # line_width=1.5,
        # xmax=140,
        # ymax=175,
        # plot_size=plot_size,
        # title='Gridded Calibration - Strong Scaling Speedup - Model Calculations',
        # show=show_instead_of_save,
        # file_name='speedup_iterations',
        # file_extension='png')
<commit_msg>efficiency plot working. started pulling out values that will be script inputs<commit_after>#!/bin/python
import sys
import os
import copy
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

def save(path, ext='png', close=True, verbose=True):
    """Save a figure from pyplot.

    Parameters
    ----------
    path : string
        The path (and filename, without the extension) to save the
        figure to.

    ext : string (default='png')
        The file extension. This must be supported by the active
        matplotlib backend (see matplotlib.backends module).  Most
        backends support 'png', 'pdf', 'ps', 'eps', and 'svg'.

    close : boolean (default=True)
        Whether to close the figure after saving.  If you want to save
        the figure multiple times (e.g., to multiple formats), you
        should NOT close it in between saves or you will have to
        re-plot it.

    verbose : boolean (default=True)
        Whether to print information about when and where the image
        has been saved.

    """

    # Extract the directory and filename from the given path
    directory = os.path.split(path)[0]
    filename = "%s.%s" % (os.path.split(path)[1], ext)
    if directory == '':
        directory = '.'

    # If the directory does not exist, create it
    if not os.path.exists(directory):
        os.makedirs(directory)

    # The final path to save to
    savepath = os.path.join(directory, filename)

    if verbose:
        print("Saving figure to '%s'..." % savepath),

    # Actually save the figure
    plt.savefig(savepath)

    # Close it
    if close:
        plt.close()

    if verbose:
        print("Done")

def read_dataframe_from_excel(
        filename,
        worksheet,
        **kwargs):
    """Reads a pandas dataframe from an Excel file
    filename: the excel file
    worksheet: worksheet name within the excel file
    The remaining named arguments are passed to pandas.read_excel
    """
    results = pd.read_excel(filename, worksheet, **kwargs)
    return results

def calculate_speedup_and_efficiency(
        rdf,
        compute_element_col_index,
        time_col_index):
    """Calculates the speedup and efficiency and
    adds them as new columns to a dataframe

    rdf: the results dataframe
    compute_element_col_index: compute elements column name
    time_col_index: column name containing the computation times

    Returns: a tuple containing the speedup and efficiency series
    """
    # get the t1 reference value
    t1 = float(rdf[rdf[compute_element_col_index] == 1][time_col_index])
    # speedup n = t1 / tn
    speedup = t1 / rdf[time_col_index]
    # efficiency n = t1 / (tn * n)
    efficiency = t1 / (rdf[time_col_index] * rdf[compute_element_col_index])

    return (speedup, efficiency)

def plot_walltime(
        series,
        colours,
        series_names,
        compute_elements,
        ymax,
        plot_size=None,
        group_width=0.8,
        xlabel='Processes',
        ylabel='Minutes',
        title='Walltime',
        file_name='walltime',
        file_extension='png',
        y_log_scale=False,
        show=False):
    """creates a bar plot as a new pyplot figure"""

    # the x locations for the groups.
    x_ind = np.arange(len(compute_elements))

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # create the individual bars
    bars_per_group = len(series)
    bar_width = group_width / bars_per_group
    bar_left = x_ind - (group_width / 2)  # centre the group on the ticks
    plot_left = x_ind[0] - group_width
    plot_right = x_ind[-1] + group_width

    for bar_values, bar_colour in zip(series, colours):
        ax.bar(
            bar_left,
            bar_values,
            width=bar_width,
            color=bar_colour,
            log=y_log_scale)
        bar_left += bar_width

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(0, ymax)
    ax.set_xticks(x_ind)
    ax.set_xticklabels(compute_elements)
    ax.legend(series_names)
    # ax.legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def plot_speedup(
        series,
        colours,
        series_names,
        compute_elements,
        xmax,
        ymax,
        line_width=1,
        plot_size=None,
        xlabel='Processes',
        ylabel='Speedup',
        title='Speedup',
        file_name='speedup',
        file_extension='png',
        show=False):
    """creates a speedup plot"""

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # define the sizes and locations of things
    x_ticks = copy.deepcopy(compute_elements)
    x_ticks.append(xmax)
    plot_left = x_ticks[0]
    plot_right = x_ticks[-1]
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(plot_left, ymax)

    # create the ideal speedup reference line
    ax.plot(
        x_ticks,
        x_ticks,
        color='red',
        linewidth=line_width,
        label='_nolegend_')

    # plot each series
    for values, colour, name in zip(series, colours, series_names):
        ax.plot(
            compute_elements,
            values,
            label=name,
            color=colour,
            linewidth=line_width,
            marker='o')

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(compute_elements)
    ax.set_xticklabels(compute_elements)
    ax.legend()
    ax.legend().get_frame().set_facecolor('white')
    # can also use get_frame().set_xy(x,y) to give a custom location

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def plot_efficiency(
        series,
        colours,
        series_names,
        compute_elements,
        xmax,
        ymax=1.2,
        line_width=1,
        plot_size=None,
        xlabel='Processes',
        ylabel='Efficiency',
        title='Efficiency',
        file_name='efficiency',
        file_extension='png',
        show=False):
    """creates an efficiency plot"""

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # define the sizes and locations of things
    x_ticks = copy.deepcopy(list(compute_elements))
    x_ticks.append(xmax)
    plot_left = x_ticks[0]
    plot_right = x_ticks[-1]
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(0, ymax)

    # create a horizontal ideal efficiency reference line
    ax.hlines(
        1,
        plot_left,
        plot_right,
        colors='red',
        linestyles='solid',
        linewidth=line_width,
        label='_nolegend_')

    # plot each series
    for values, colour, name in zip(series, colours, series_names):
        ax.plot(
            compute_elements,
            values,
            label=name,
            color=colour,
            linewidth=line_width,
            marker='o')

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(compute_elements)
    ax.set_xticklabels(compute_elements)
    ax.legend(shadow=True)
    ax.legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

if __name__ == '__main__':

    # create plots in a 4:3 aspect ratio
    # matplotlib works in inches
    plot_width = 10
    plot_size = (plot_width, plot_width * 3 / 4)

    # show_instead_of_save = False
    show_instead_of_save = True

    # some values that should be inputs to this file
    plot_titles = { 'walltime':'ODPS - Walltime', 'speedup':'ODPS - Speedup', 'efficiency':'ODPS - Strong Scaling Efficiency'}
    compute_element_name = 'Threads'
    walltime_units = 'Seconds'

    # CSIRO colours
    colours = [
        '#00a9ce',  # midday blue
        '#00616c',  # midnight blue
        '#78be20',  # light forest
        '#DF1995',  # fuschia
        '#E87722',  # orange
        '#E4002B',  # vermillion
        '#FFB81C',  # gold
        '#6D2077',  # plum
        '#1E22AA',  # blueberry
    ]

    results = read_dataframe_from_excel(
        'results.xls',
        worksheet='results',
        usecols=[
            'group',
            'compute_elements',
            'walltime'])

    # I don't know an equivalent to R's completecases, but this works for now
    results = results[results['group'].notnull() & (results.walltime > 0)]

    # create new dataframes for each group of results
    groups = results.group.unique()
    group_dataframes = {}
    for group in groups:
        df = results[(results.group == group)]
        if len(df):
            group_dataframes[group] = df.sort('compute_elements', ascending=True)

    # calculate speedup and efficiency based on the total walltime, and on
    # the average time per iteration (which excludes a lot of the serial
    # code, and gives a closer look at just the model calculations)
    for data in group_dataframes.values():
        speedup, efficiency = calculate_speedup_and_efficiency(
            data,
            'compute_elements',
            'walltime')
        data['speedup'] = speedup
        data['efficiency'] = efficiency

    # extract the compute element names for grouping and labelling the charts
    compute_elements = np.sort(results.compute_elements.unique())

    # build the collections of values for the charts
    walltimes = []
    efficiencies = []
    speedups = []
    series_names = []
    max_walltime = 0
    for name, data in group_dataframes.items():
        walltimes.append(data.walltime)
        efficiencies.append(data.efficiency)
        speedups.append(data.speedup)
        series_names.append(name)
        local_max_walltime = max(data.walltime)
        if local_max_walltime > max_walltime:
            max_walltime = local_max_walltime

    # finally make the plots

    # plot_walltime(
        # walltimes,
        # colours,
        # series_names,
        # compute_elements,
        # ymax=max_walltime + 50,
        # group_width=0.83,
        # plot_size=plot_size,
        # title='ODPS Walltime',
        # xlabel=compute_element_name,
        # ylabel=walltime_units,
        # show=show_instead_of_save,
        # file_name='odps-walltime',
        # file_extension='png')

    plot_efficiency(
        efficiencies,
        colours,
        series_names,
        compute_elements,
        line_width=1.5,
        xmax=compute_elements.max() * 1.1,
        ymax=1.3,
        plot_size=plot_size,
        xlabel=compute_element_name,
        title=plot_titles['efficiency'],
        show=show_instead_of_save,
        file_name='odps-efficiency',
        file_extension='png')

    # plot_speedup(
        # speedups,
        # colours,
        # series_names,
        # compute_elements,
        # line_width=1.5,
        # xmax=140,
        # ymax=140,
        # plot_size=plot_size,
        # title='Gridded Calibration - Strong Scaling Speedup',
        # show=show_instead_of_save,
        # file_name='speedup',
        # file_extension='png')
<|endoftext|>


<commit_before>
#from . import Chart
def make_chart(df,**kwargs):

    c = Chart(**kwargs)

    elem_series = {
            'name': '',
            'type': kwargs['type'],
            'data': []}

    x = kwargs['x']
    y = kwargs.get('y')
    category = kwargs['category']

    def round_list(arr):
        try:
            return arr.values.round(3).tolist()
        except TypeError:
            return arr.unique().tolist()

    if kwargs['type'] in ['bar','line']:

        def insert_series(df,cat=None):
            """Return data series based on dataframe"""
            c._option['xAxis']['data'] = round_list(df[x].drop_duplicates())

            data = (df[x].value_counts()
                        if y is None else
                        df.groupby(x)[y].aggregate(kwargs['how']))

            series = deepcopy(elem_series)
            series['data'] = round_list(data)
            series['type'] = kwargs['type']

            if kwargs['stacked'] == True:
                series['stack'] = category

            series['name'] = cat if cat else x
            c._option['series'].append(series)

        if category:
            #Iterate and append Data for every category
            for cat, subset in df.groupby(category):
                insert_series(subset,cat=cat)
                c._option['legend']['data'].append(cat)
        else:
            insert_series(df)

        if (category and
            kwargs['type'] == 'line' and
            kwargs['area'] == True and
            kwargs['stacked']== True):
            for e in c._option['series']:
                e['areaStyle']= {'normal': {}}

    elif kwargs['type'] == 'hist':
        kwargs['type'] = 'bar'

        def insert_series(df,cat=None):

            y_val,x_val = np.histogram(df[x],
                                       bins=kwargs['bins'],
                                       normed=kwargs['normed'])
            data = pd.Series(y_val)
            bins = x_val.astype(int).tolist()
            c._option['xAxis']['data'] = bins

            series = deepcopy(elem_series)
            series['data'] = round_list(data)
            series['type'] = kwargs['type']
            series['name'] = cat if cat else x

            if kwargs['stacked'] == True:
                series['stack'] = category

            c._option['series'].append(series)

        if category:
            for cat,subset in df.groupby(category):
                insert_series(subset,cat=cat)
                c._option['legend']['data'].append(cat)
        else:
            insert_series(df)
#             c._option['xAxis']['data'] =  get_hist_series(df,x,**kwargs)

    elif kwargs['type'] == 'scatter':

        c._option['xAxis'] = {'type': 'value',
                              'name': x,
                              'max': int(df[x].max())}

        c._option['yAxis'] = {'type': 'value',
                              'name': y,
                              'max': int(df[y].max())}

        cols = [x,y]

        c._option['visualMap'] = []
        visual_map_template = {'show': False,
                               'min': 0,
                               'max': 999,
                               'inRange': {}}

        size = kwargs['size']
        if size is not None:
            vmap_size = deepcopy(visual_map_template)
            vmap_size['min'] = df[size].min()
            vmap_size['max'] = df[size].max()
            vmap_size['inRange']['symbolSize'] = [6,60]
            c._option['visualMap'].append(vmap_size)
            cols.append(size)

        #TODO: Fix Saturate
#          saturate = kwargs['saturate']
#         if saturate is not None:
#             vmap_saturate = deepcopy(visual_map_template)
#             vmap_saturate['min'] = float(df[saturate].min())
#             vmap_saturate['max'] = float(df[saturate].max())
#             vmap_saturate['inRange']['colorLightness'] = [1,0.5]
#             c._option['visualMap'].append(vmap_saturate)
#             cols.append(saturate)

        def insert_series(df,cat=None):

            data = df[cols]
            series = deepcopy(elem_series)
            series['data'] = round_list(data)
            series['type'] = kwargs['type']
            series['name'] = cat if cat else x
            c._option['series'].append(series)

        if category:
            for cat,subset in df.groupby(category):
                insert_series(subset,cat=cat)
                c._option['legend']['data'].append(cat)
        else:
            insert_series(df)

    return c

def bar(df,x,y=None,category=None,how='count',stacked=False,**kwargs):

    kwargs['x'] = x
    kwargs['y'] = y
    kwargs['category'] = category
    kwargs['how'] = how
    kwargs['type'] = 'bar'
    kwargs['stacked'] = stacked

    return make_chart(df,**kwargs)

def line(df,x,y=None,category=None,how='count',stacked=False,area=False,**kwargs):

    kwargs['x'] = x
    kwargs['y'] = y
    kwargs['category'] = category
    kwargs['how'] = how
    kwargs['type'] = 'line'
    kwargs['stacked'] = stacked
    kwargs['area'] = area

    return make_chart(df,**kwargs)

def hist(df,x,category=None,bins=10,normed=False,stacked=False,**kwargs):

    kwargs['x'] = x
    kwargs['category'] = category
    kwargs['bins'] = bins
    kwargs['type'] = 'hist'
    kwargs['normed'] = normed
    kwargs['stacked'] = stacked

    return make_chart(df,**kwargs)

def scatter(df,x,y,size=None,category=None,**kwargs):

    kwargs['x'] = x
    kwargs['y'] = y
    kwargs['category'] = category
    kwargs['size'] = size
    #kwargs['saturate'] = saturate #TODO: Fix saturate
    kwargs['type'] = 'scatter'

    return make_chart(df,**kwargs)<commit_msg>Refactor series<commit_after>
#from . import Chart

def bar_line_hist_condition(c,category=None,**kwargs):
    """Provide stacked,annotate, area for bar line hist"""
    series = c._option['series']

    d_annotate = {'normal':{'show':True,
                            'position':'top'}}

    if category and kwargs['stacked'] == True:
        for s in series:
            s['stack'] = category

            if kwargs['type'] == 'line' and kwargs['area'] == True:
                s['areaStyle'] = {'normal': {}}

            if kwargs['annotate'] == 'all':
                s['label'] = d_annotate

    if kwargs['annotate'] == 'top':
        series[-1]['label'] = d_annotate

def round_list(arr):
    try:
        return arr.values.round(3).tolist()
    except TypeError:
        return arr.unique().tolist()

def make_chart(df,**kwargs):

    def insert_series_on(f):

        if category:
            #Iterate and append Data for every category
            for cat, subset in df.groupby(category):
                cat = str(cat)
                insert_data_to_series(f,subset,cat)
                c._option['legend']['data'].append(cat)
        else:
            insert_data_to_series(f,df)

    def insert_data_to_series(f,df,cat=None):
        data = f(df)

        series = deepcopy(elem_series)
        series['data'] = round_list(data)
        series['type'] = kwargs['type']
        series['name'] = cat if cat else x
        c._option['series'].append(series)

    c = Chart(**kwargs)

    elem_series = {
            'name': '',
            'type': kwargs['type'],
            'data': []}

    x = kwargs['x']
    y = kwargs.get('y')
    category = kwargs['category']

    if kwargs['type'] in ['bar','line']:

        def get_bar_line_data(df):
            c._option['xAxis']['data'] = round_list(df[x].drop_duplicates())

            data = (df[x].value_counts()
                    if y is None else
                    df.groupby(x)[y].aggregate(kwargs['how']))

            return data

        insert_series_on(get_bar_line_data)
        bar_line_hist_condition(c,**kwargs)

    elif kwargs['type'] == 'hist':
        kwargs['type'] = 'bar'

        def get_hist_data(df):

            y_val,x_val = np.histogram(df[x],
                                       bins=kwargs['bins'],
                                       normed=kwargs['normed'])
            data = pd.Series(y_val)
            bins = x_val.astype(int).tolist()
            c._option['xAxis']['data'] = bins

            return data

        insert_series_on(get_hist_data)
        bar_line_hist_condition(c,**kwargs)

    elif kwargs['type'] == 'scatter':

        c._option['xAxis'] = {'type': 'value',
                              'name': x,
                              'max': int(df[x].max())}
        c._option['yAxis'] = {'type': 'value',
                              'name': y,
                              'max': int(df[y].max())}
        c._option['visualMap'] = []

        visual_map_template = {'show': False,
                               'min': 0,
                               'max': 999,
                               'inRange': {}}

        cols = [x,y]
        size = kwargs['size']
        if size is not None:
            vmap_size = deepcopy(visual_map_template)
            vmap_size['min'] = df[size].min()
            vmap_size['max'] = df[size].max()
            vmap_size['inRange']['symbolSize'] = [6,60]
            c._option['visualMap'].append(vmap_size)
            cols.append(size)

        #TODO: Fix Saturate
#          saturate = kwargs['saturate']
#         if saturate is not None:
#             vmap_saturate = deepcopy(visual_map_template)
#             vmap_saturate['min'] = float(df[saturate].min())
#             vmap_saturate['max'] = float(df[saturate].max())
#             vmap_saturate['inRange']['colorLightness'] = [1,0.5]
#             c._option['visualMap'].append(vmap_saturate)
#             cols.append(saturate)

        def get_scatter_data(df):

            data = df[cols]
            return data

        insert_series_on(get_scatter_data)

    return c

def bar(df,x,y=None,category=None,how='count',stacked=False,
        annotate=None,**kwargs):

    kwargs['x'] = x
    kwargs['y'] = y
    kwargs['category'] = category
    kwargs['how'] = how
    kwargs['type'] = 'bar'
    kwargs['stacked'] = stacked
    kwargs['annotate'] = 'top' if annotate == True else annotate

    return make_chart(df,**kwargs)

def line(df,x,y=None,category=None,how='count',stacked=False,area=False,
         annotate=None,**kwargs):

    kwargs['x'] = x
    kwargs['y'] = y
    kwargs['category'] = category
    kwargs['how'] = how
    kwargs['type'] = 'line'
    kwargs['stacked'] = stacked
    kwargs['area'] = area
    kwargs['annotate'] = 'top' if annotate == True else annotate

    return make_chart(df,**kwargs)

def hist(df,x,category=None,bins=10,normed=False,stacked=False,
         annotate=None,**kwargs):

    kwargs['x'] = x
    kwargs['category'] = category
    kwargs['bins'] = bins
    kwargs['type'] = 'hist'
    kwargs['normed'] = normed
    kwargs['stacked'] = stacked
    kwargs['annotate'] = 'top' if annotate == True else annotate

    return make_chart(df,**kwargs)

def scatter(df,x,y,size=None,category=None,**kwargs):

    kwargs['x'] = x
    kwargs['y'] = y
    kwargs['category'] = category
    kwargs['size'] = size
    #kwargs['saturate'] = saturate #TODO: Fix saturate
    kwargs['type'] = 'scatter'

    return make_chart(df,**kwargs)<|endoftext|>


<commit_before>#   Foremast - Pipeline Tooling
#
#   Copyright 2016 Gogo, LLC
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import logging
import os
import subprocess

from ..utils import get_properties, get_details
from ..exceptions import S3ArtifactNotFound

LOG = logging.getLogger(__name__)

class S3Deployment(object):
    """Handles uploading artifact to S3 and S3 deployment strategies"""

    def __init__(self, app, env, region, prop_path, artifact_path, artifact_version):
        """S3 deployment object.

        Args:
            app (str): Application name
            env (str): Environment/Account
            region (str): AWS Region
            prop_path (str): Path of environment property file
            artifact_path (str): Path to tar.gz artifact
        """
        self.app_name = app
        self.env = env
        self.region = region
        self.artifact_path = artifact_path
        self.version = artifact_version
        generated = get_details(app=app, env=env)
        self.bucket = generated.s3_app_bucket()
        self.properties = get_properties(prop_path)
        self.s3props = self.properties[self.env]['s3']
        self.s3path = self.s3props['path'].lstrip('/')
        self.s3_version_uri = ''
        self.s3_latest_uri = ''
        self.setup_pathing()

    def setup_pathing(self):
        """Formats pathing for S3 deployments"""
        path_format = "{}/{}/{}"
        s3_format = "s3://{}"
        version_path = path_format.format(self.bucket, self.s3path, self.version).replace('//', '/')
        latest_path = path_format.format(self.bucket, self.s3path, "LATEST").replace('//', '/')
        self.s3_version_uri = s3_format.format(version_path)
        self.s3_latest_uri = s3_format.format(latest_path)

    def upload_artifacts(self):
        """Uploads the artifacts to S3 and copies to LATEST depending on strategy"""
        deploy_strategy = self.properties[self.env]["deploy_strategy"]
        if deploy_strategy == "highlander":
            self._upload_artifacts_to_version()
            self._sync_to_latest()
        elif deploy_strategy == "redblack":
            self._upload_artifacts_to_version()
        else:
            raise NotImplemented

    def promote_artifacts(self):
        """Promotes artifact version to LATEST"""
        self._sync_to_latest()

    def _upload_artifacts_to_version(self):
        """Recursively uploads a directory and all files and subdirectories to S3"""
        if not os.listdir(self.artifact_path) or not self.artifact_path:
            raise S3ArtifactNotFound

        cmd = 'aws s3 sync {} {} --delete --profile {}'.format(self.artifact_path, self.s3_version_uri, self.env)
        p = subprocess.run(cmd, check=True, shell=True, stdout=subprocess.PIPE)
        LOG.debug("Upload Command Ouput: %s", p.stdout)
        LOG.info("Uploaded artifacts to %s bucket", self.bucket)

    def _sync_to_latest(self):
        """Uses AWS CLI to sync versioned directory to LATEST directory in S3"""
        cmd = 'aws s3 sync {} {} --delete --profile {}'.format(self.s3_version_uri, self.s3_latest_uri, self.env)
        p = subprocess.run(cmd, check=True, shell=True, stdout=subprocess.PIPE)
        LOG.debug("Sync to latest command output: %s", p.stdout)
        LOG.info("Copied version %s to LATEST", self.version)
<commit_msg>added '--exact-timestamps' to s3 sync<commit_after>#   Foremast - Pipeline Tooling
#
#   Copyright 2016 Gogo, LLC
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import logging
import os
import subprocess

from ..utils import get_properties, get_details
from ..exceptions import S3ArtifactNotFound

LOG = logging.getLogger(__name__)

class S3Deployment(object):
    """Handles uploading artifact to S3 and S3 deployment strategies"""

    def __init__(self, app, env, region, prop_path, artifact_path, artifact_version):
        """S3 deployment object.

        Args:
            app (str): Application name
            env (str): Environment/Account
            region (str): AWS Region
            prop_path (str): Path of environment property file
            artifact_path (str): Path to tar.gz artifact
        """
        self.app_name = app
        self.env = env
        self.region = region
        self.artifact_path = artifact_path
        self.version = artifact_version
        generated = get_details(app=app, env=env)
        self.bucket = generated.s3_app_bucket()
        self.properties = get_properties(prop_path)
        self.s3props = self.properties[self.env]['s3']
        self.s3path = self.s3props['path'].lstrip('/')
        self.s3_version_uri = ''
        self.s3_latest_uri = ''
        self.setup_pathing()

    def setup_pathing(self):
        """Formats pathing for S3 deployments"""
        path_format = "{}/{}/{}"
        s3_format = "s3://{}"
        version_path = path_format.format(self.bucket, self.s3path, self.version).replace('//', '/')
        latest_path = path_format.format(self.bucket, self.s3path, "LATEST").replace('//', '/')
        self.s3_version_uri = s3_format.format(version_path)
        self.s3_latest_uri = s3_format.format(latest_path)

    def upload_artifacts(self):
        """Uploads the artifacts to S3 and copies to LATEST depending on strategy"""
        deploy_strategy = self.properties[self.env]["deploy_strategy"]
        if deploy_strategy == "highlander":
            self._upload_artifacts_to_version()
            self._sync_to_latest()
        elif deploy_strategy == "redblack":
            self._upload_artifacts_to_version()
        else:
            raise NotImplemented

    def promote_artifacts(self):
        """Promotes artifact version to LATEST"""
        self._sync_to_latest()

    def _upload_artifacts_to_version(self):
        """Recursively uploads a directory and all files and subdirectories to S3"""
        if not os.listdir(self.artifact_path) or not self.artifact_path:
            raise S3ArtifactNotFound

        cmd = 'aws s3 sync {} {} --delete --exact-timestamps --profile {}'.format(self.artifact_path, self.s3_version_uri, self.env)
        p = subprocess.run(cmd, check=True, shell=True, stdout=subprocess.PIPE)
        LOG.debug("Upload Command Ouput: %s", p.stdout)
        LOG.info("Uploaded artifacts to %s bucket", self.bucket)

    def _sync_to_latest(self):
        """Uses AWS CLI to sync versioned directory to LATEST directory in S3"""
        cmd = 'aws s3 sync {} {} --delete --exact-timestamps --profile {}'.format(self.s3_version_uri, self.s3_latest_uri, self.env)
        p = subprocess.run(cmd, check=True, shell=True, stdout=subprocess.PIPE)
        LOG.debug("Sync to latest command output: %s", p.stdout)
        LOG.info("Copied version %s to LATEST", self.version)
<|endoftext|>


<commit_before>#  Copyright 2008-2010 Nokia Siemens Networks Oyj
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import sys

from robot import utils
from robot.errors import FrameworkError

from loggerhelper import IsLogged

class CommandLineMonitor:

    def __init__(self, width=78, colors=True):
        self._width = width
        self._colors = colors
        self._running_suites = 0
        self._is_logged = IsLogged('WARN')

    def start_suite(self, suite):
        if not self._running_suites:
            self._write_separator('=')
        self._write_info(suite.longname, suite.doc, start_suite=True)
        self._write_separator('=')
        self._running_suites += 1

    def end_suite(self, suite):
        self._write_info(suite.longname, suite.doc)
        self._write_status(suite.status)
        self._write_message(suite.get_full_message())
        self._write_separator('=')
        self._running_suites -= 1

    def start_test(self, test):
        self._write_info(test.name, test.doc)

    def end_test(self, test):
        self._write_status(test.status)
        self._write_message(test.message)
        self._write_separator('-')

    def output_file(self, name, path):
        # called by LOGGER
        if not self._running_suites:  # ignore split output files
            self._write('%s %s' % ((name+':').ljust(8), path))

    def message(self, msg):
        # called by LOGGER
        if self._is_logged(msg.level):
            message = '[ %s ] %s' % (self._status_text(msg.level), msg.message)
            self._write(message, stream=sys.__stderr__)

    def _status_text(self, text):
        return StatusText(text, self._colors)

    def _write(self, message, newline=True, stream=sys.__stdout__):
        if newline:
            message += '\n'
        stream.write(utils.encode_output(message).replace('\t', ' '*8))
        stream.flush()

    def _write_info(self, name, doc, start_suite=False):
        maxwidth = self._width
        if not start_suite:
            maxwidth -= len(' | PASS |')
        info = self._get_info(name, doc, maxwidth)
        self._write(info, newline=start_suite)

    def _get_info(self, name, doc, maxwidth):
        if utils.get_console_length(name) > maxwidth:
            return utils.pad_console_length(name, maxwidth, cut_left=True)
        if doc == '':
            return utils.pad_console_length(name, maxwidth)
        info = '%s :: %s' % (name, doc.splitlines()[0])
        return utils.pad_console_length(info, maxwidth)

    def _write_status(self, status):
        self._write(' | %s |' % self._status_text(status))

    def _write_message(self, message):
        if message:
            self._write(message.strip())

    def _write_separator(self, sep_char):
        self._write(sep_char * self._width)

def StatusText(msg, colors=True):
    if colors:
        return HiglightedStatusText(msg)
    return PlainStatusText(msg)

class PlainStatusText:

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

    def __str__(self):
        return self._msg

class HiglightedStatusText(PlainStatusText):
    ANSI_RED    = '\033[31m'
    ANSI_GREEN  = '\033[32m'
    ANSI_YELLOW = '\033[33m'
    ANSI_RESET  = '\033[0m'

    def __str__(self):
        color = self._get_highlight_color(self._msg)
        reset = color != '' and self.ANSI_RESET or ''
        return color + self._msg + reset

    def _get_highlight_color(self, text):
        if text in ['FAIL','ERROR']:
            return self.ANSI_RED
        elif text == 'WARN':
            return self.ANSI_YELLOW
        elif text == 'PASS':
            return self.ANSI_GREEN
        raise FrameworkError<commit_msg>monitor: Move writing of the possibly highlihgted texts into PlainStatusText baseclass<commit_after>#  Copyright 2008-2010 Nokia Siemens Networks Oyj
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import sys

from robot import utils
from robot.errors import FrameworkError

from loggerhelper import IsLogged

class CommandLineMonitor:

    def __init__(self, width=78, colors=True):
        self._width = width
        self._colors = colors
        self._running_suites = 0
        self._is_logged = IsLogged('WARN')

    def start_suite(self, suite):
        if not self._running_suites:
            self._write_separator('=')
        self._write_info(suite.longname, suite.doc, start_suite=True)
        self._write_separator('=')
        self._running_suites += 1

    def end_suite(self, suite):
        self._write_info(suite.longname, suite.doc)
        self._write_status(suite.status)
        self._write_message(suite.get_full_message())
        self._write_separator('=')
        self._running_suites -= 1

    def start_test(self, test):
        self._write_info(test.name, test.doc)

    def end_test(self, test):
        self._write_status(test.status)
        self._write_message(test.message)
        self._write_separator('-')

    def output_file(self, name, path):
        # called by LOGGER
        if not self._running_suites:  # ignore split output files
            self._write('%s %s' % ((name+':').ljust(8), path))

    def message(self, msg):
        # called by LOGGER
        if self._is_logged(msg.level):
            self._status_text(msg.level).write_msg(msg.message)

    def _status_text(self, text):
        return StatusText(text, self._colors)

    def _write(self, message, newline=True, stream=sys.__stdout__):
        if newline:
            message += '\n'
        stream.write(utils.encode_output(message).replace('\t', ' '*8))
        stream.flush()

    def _write_info(self, name, doc, start_suite=False):
        maxwidth = self._width
        if not start_suite:
            maxwidth -= len(' | PASS |')
        info = self._get_info(name, doc, maxwidth)
        self._write(info, newline=start_suite)

    def _get_info(self, name, doc, maxwidth):
        if utils.get_console_length(name) > maxwidth:
            return utils.pad_console_length(name, maxwidth, cut_left=True)
        if doc == '':
            return utils.pad_console_length(name, maxwidth)
        info = '%s :: %s' % (name, doc.splitlines()[0])
        return utils.pad_console_length(info, maxwidth)

    def _write_status(self, status):
        self._status_text(status).write_status()

    def _write_message(self, message):
        if message:
            self._write(message.strip())

    def _write_separator(self, sep_char):
        self._write(sep_char * self._width)

def StatusText(msg, colors=True):
    if colors:
        return HiglightedStatusText(msg)
    return PlainStatusText(msg)

class PlainStatusText:

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

    def __str__(self):
        return self._msg

    def write_status(self, stream=sys.__stdout__):
        self.write(' | %s |' % self, stream)

    def write_msg(self, msg):
        self.write('[ %s ] %s' % (self, msg), stream=sys.__stderr__)

    def write(self, message, newline=True, stream=sys.__stdout__):
        if newline:
            message += '\n'
        stream.write(utils.encode_output(message).replace('\t', ' '*8))
        stream.flush()

class HiglightedStatusText(PlainStatusText):
    ANSI_RED    = '\033[31m'
    ANSI_GREEN  = '\033[32m'
    ANSI_YELLOW = '\033[33m'
    ANSI_RESET  = '\033[0m'

    def __str__(self):
        color = self._get_highlight_color(self._msg)
        reset = color != '' and self.ANSI_RESET or ''
        return color + self._msg + reset

    def _get_highlight_color(self, text):
        if text in ['FAIL','ERROR']:
            return self.ANSI_RED
        elif text == 'WARN':
            return self.ANSI_YELLOW
        elif text == 'PASS':
            return self.ANSI_GREEN
        raise FrameworkError
<|endoftext|>


<commit_before>#
# Copyright (c) 2011 Red Hat, Inc.
#
# This software is licensed to you under the GNU Lesser General Public
# License as published by the Free Software Foundation; either version
# 2 of the License (LGPLv2) or (at your option) any later version.
# There is NO WARRANTY for this software, express or implied,
# including the implied warranties of MERCHANTABILITY,
# NON-INFRINGEMENT, or FITNESS FOR A PARTICULAR PURPOSE. You should
# have received a copy of LGPLv2 along with this software; if not, see
# http://www.gnu.org/licenses/old-licenses/lgpl-2.0.txt.
#
# Jeff Ortel <jortel@redhat.com>
#

"""
Thread Pool classes.
"""

from uuid import uuid4
from threading import Thread
from Queue import Queue
from logging import getLogger

log = getLogger(__name__)

class Worker(Thread):
    """
    Pool (worker) thread.
    :ivar queue: A thread pool worker.
    :type queue: Queue
    """

    def __init__(self, worker_id, backlog=100):
        """
        :param worker_id: The worker id in the pool.
        :type worker_id: int
        :param backlog: Limits the number of calls queued.
        :type backlog: int
        """
        name = 'worker-%d' % worker_id
        Thread.__init__(self, name=name)
        self.queue = Queue(backlog)
        self.setDaemon(True)

    def run(self):
        """
        Main run loop; processes input queue.
        """
        while True:
            call = self.queue.get()
            if not call:
                # termination requested
                return
            try:
                call()
            except Exception:
                log.exception(str(call))

    def put(self, call):
        """
        Enqueue a call.
        :param call: A call to queue.
        :type call: Call
        """
        self.queue.put(call)

    def backlog(self):
        """
        Get the number of call already queued to this worker.
        :return: The number of queued calls.
        :rtype: int
        """
        return self.queue.qsize()

class Call:
    """
    A call to be executed by the thread pool.
    :ivar id: The unique call ID.
    :type id: str
    :ivar fn: The function/method to be executed.
    :type fn: callable
    :ivar args: The list of args passed to the callable.
    :type args: list
    :ivar kwargs: The list of keyword args passed to the callable.
    :type kwargs: dict
    """

    def __init__(self, call_id, fn, args=None, kwargs=None):
        """
        :param call_id: The unique call ID.
        :type call_id: str
        :param fn: The function/method to be executed.
        :type fn: callable
        :param args: The list of args passed to the callable.
        :type args: tuple
        :param kwargs: The list of keyword args passed to the callable.
        :type kwargs: dict
        """
        self.id = call_id
        self.fn = fn
        self.args = args or []
        self.kwargs = kwargs or {}

    def __call__(self):
        """
        Execute the call.
        """
        try:
            self.fn(*self.args, **self.kwargs)
        except Exception:
            log.exception(self.id)

    def __str__(self):
        s = list()
        s.append('call: ')
        s.append(str(self.fn))
        s.append('(')
        s.append(str(self.args))
        s.append(', ')
        s.append(str(self.kwargs))
        s.append(')')
        return ''.join(s)

class ThreadPool:
    """
    A load distributed thread pool.
    :ivar capacity: The min # of workers.
    :type capacity: int
    """

    def __init__(self, capacity=1):
        """
        :param capacity: The # of workers.
        :type capacity: int
        """
        self.capacity = capacity
        self.threads = []
        for x in range(capacity):
            self.__add()

    def run(self, fn, *args, **kwargs):
        """
        Schedule a call.
        Convenience method for scheduling.
        :param fn: A function/method to execute.
        :type fn: callable
        :param args: The args passed to fn()
        :type args: tuple
        :param kwargs: The keyword args passed fn()
        :type kwargs: dict
        :return The call ID.
        :rtype str
        """
        call_id = uuid4()
        call = Call(call_id, fn, args, kwargs)
        return self.schedule(call)

    def schedule(self, call):
        """
        Schedule a call.
        :param call: A call to schedule for execution.
        :param call: Call
        :return: The call ID.
        :rtype: str
        """
        pool = [(t.backlog(), t) for t in self.threads]
        pool.sort()
        backlog, worker = pool[0]
        worker.put(call)

    def shutdown(self):
        """
        Shutdown the pool.
        Terminate and join all workers.
        """
        # send stop request
        for t in self.threads:
            t.put(0)
        for t in self.threads:
            t.join()

    def __add(self):
        """
        Add a thread to the pool.
        """
        n = len(self.threads)
        thread = Worker(n)
        self.threads.append(thread)
        thread.start()

    def __len__(self):
        return len(self.threads)

    def __repr__(self):
        s = list()
        s.append('pool: capacity=%d' % self.capacity)
        for t in self.threads:
            s.append('worker: %s backlog: %d' % (t.name, t.backlog()))
        return '\n'.join(s)

class Direct:
    """
    Call ignored (trashed).
    """

    def run(self, fn, *args, **options):
        """
        Run request.
        :param fn: A function/method to execute.
        :type fn: callable
        :param args: The args passed to fn()
        :type args: list
        :param options: The keyword args passed fn()
        :type options: dict
        """
        return fn(*args, **options)
<commit_msg>Fix thread pool worker selection.<commit_after>#
# Copyright (c) 2011 Red Hat, Inc.
#
# This software is licensed to you under the GNU Lesser General Public
# License as published by the Free Software Foundation; either version
# 2 of the License (LGPLv2) or (at your option) any later version.
# There is NO WARRANTY for this software, express or implied,
# including the implied warranties of MERCHANTABILITY,
# NON-INFRINGEMENT, or FITNESS FOR A PARTICULAR PURPOSE. You should
# have received a copy of LGPLv2 along with this software; if not, see
# http://www.gnu.org/licenses/old-licenses/lgpl-2.0.txt.
#
# Jeff Ortel <jortel@redhat.com>
#

"""
Thread Pool classes.
"""

from uuid import uuid4
from threading import Thread
from Queue import Queue
from logging import getLogger

log = getLogger(__name__)

class Worker(Thread):
    """
    Pool (worker) thread.
    The *busy* marker is queued to ensure that the backlog
    for busy workers is longer which supports better work distribution.
    :ivar queue: A thread pool worker.
    :type queue: Queue
    """

    def __init__(self, worker_id, backlog=100):
        """
        :param worker_id: The worker id in the pool.
        :type worker_id: int
        :param backlog: Limits the number of calls queued.
        :type backlog: int
        """
        name = 'worker-%d' % worker_id
        Thread.__init__(self, name=name)
        self.queue = Queue(backlog)
        self.setDaemon(True)

    def run(self):
        """
        Main run loop; processes input queue.
        """
        while True:
            call = self.queue.get()
            if call == 1:
                # # busy
                continue
            if not call:
                # termination requested
                return
            try:
                call()
            except Exception:
                log.exception(str(call))

    def put(self, call):
        """
        Enqueue a call.
        :param call: A call to queue.
        :type call: Call
        """
        self.queue.put(call)
        self.queue.put(1)  # busy

    def backlog(self):
        """
        Get the number of call already queued to this worker.
        :return: The number of queued calls.
        :rtype: int
        """
        return self.queue.qsize()

class Call:
    """
    A call to be executed by the thread pool.
    :ivar id: The unique call ID.
    :type id: str
    :ivar fn: The function/method to be executed.
    :type fn: callable
    :ivar args: The list of args passed to the callable.
    :type args: list
    :ivar kwargs: The list of keyword args passed to the callable.
    :type kwargs: dict
    """

    def __init__(self, call_id, fn, args=None, kwargs=None):
        """
        :param call_id: The unique call ID.
        :type call_id: str
        :param fn: The function/method to be executed.
        :type fn: callable
        :param args: The list of args passed to the callable.
        :type args: tuple
        :param kwargs: The list of keyword args passed to the callable.
        :type kwargs: dict
        """
        self.id = call_id
        self.fn = fn
        self.args = args or []
        self.kwargs = kwargs or {}

    def __call__(self):
        """
        Execute the call.
        """
        try:
            self.fn(*self.args, **self.kwargs)
        except Exception:
            log.exception(self.id)

    def __str__(self):
        s = list()
        s.append('call: ')
        s.append(str(self.fn))
        s.append('(')
        s.append(str(self.args))
        s.append(', ')
        s.append(str(self.kwargs))
        s.append(')')
        return ''.join(s)

class ThreadPool:
    """
    A load distributed thread pool.
    :ivar capacity: The min # of workers.
    :type capacity: int
    """

    def __init__(self, capacity=1):
        """
        :param capacity: The # of workers.
        :type capacity: int
        """
        self.capacity = capacity
        self.threads = []
        for x in range(capacity):
            self.__add()

    def run(self, fn, *args, **kwargs):
        """
        Schedule a call.
        Convenience method for scheduling.
        :param fn: A function/method to execute.
        :type fn: callable
        :param args: The args passed to fn()
        :type args: tuple
        :param kwargs: The keyword args passed fn()
        :type kwargs: dict
        :return The call ID.
        :rtype str
        """
        call_id = uuid4()
        call = Call(call_id, fn, args, kwargs)
        return self.schedule(call)

    def schedule(self, call):
        """
        Schedule a call.
        :param call: A call to schedule for execution.
        :param call: Call
        :return: The call ID.
        :rtype: str
        """
        pool = [(t.backlog(), t) for t in self.threads]
        pool.sort()
        backlog, worker = pool[0]
        worker.put(call)

    def shutdown(self):
        """
        Shutdown the pool.
        Terminate and join all workers.
        """
        # send stop request
        for t in self.threads:
            t.put(0)
        for t in self.threads:
            t.join()

    def __add(self):
        """
        Add a thread to the pool.
        """
        n = len(self.threads)
        thread = Worker(n)
        self.threads.append(thread)
        thread.start()

    def __len__(self):
        return len(self.threads)

    def __repr__(self):
        s = list()
        s.append('pool: capacity=%d' % self.capacity)
        for t in self.threads:
            s.append('worker: %s backlog: %d' % (t.name, t.backlog()))
        return '\n'.join(s)

class Direct:
    """
    Call ignored (trashed).
    """

    def run(self, fn, *args, **options):
        """
        Run request.
        :param fn: A function/method to execute.
        :type fn: callable
        :param args: The args passed to fn()
        :type args: list
        :param options: The keyword args passed fn()
        :type options: dict
        """
        return fn(*args, **options)
<|endoftext|>