reflow/reflowctl/solder.py

import os
import os.path

import xml.etree.ElementTree as etree

from PyQt4 import QtGui, QtCore

from numpy import arange, sin, pi, array, linspace, arange

from temp_level import TempLevel, set_colors
from edge_widget import Edge
from control_widgets import AddRemoveWidget


def getTemperature():
    return 20.


class Solder(QtCore.QObject):

    log_message = QtCore.pyqtSignal(str)

    def __init__(self, filename, name=str(), description=str(), parent=None):
        super(Solder, self).__init__(parent)
        self.changed = False
        self.filename = filename
        self.name = name
        self.description = description
        self.temp_levels = list()
        self.edges = list()


    def __unicode__(self):
        return unicode(self.name)

    def __str__(self):
        return self.name

    def add_temp_level(self, name, temp, is_env):
        s = TempLevel(name, temp, is_env)
        self.temp_levels.append(s)
        return s

    def get_temp_level_by_name(self, name):
        assert isinstance(name, basestring)
        for i in self.temp_levels:
            if i.name == name:
                return i
        return None

    def calc_profile(self):
        self.log = list()
        x = list()
        y = list()
        duration_points = dict()
        rate_points = dict()
        time = 0

        def calc(edge):
            if edge.duration:
                return time + edge.duration
            elif edge.rate:
                return time + (edge.to_tl.temp - edge.from_tl.temp) / edge.rate
            else:
                raise Exception("edge %r has neither duration nor rate set" % edge)

        for _edge in self.edges:
            x.append(float(time))
            y.append(float(_edge.from_tl.temp))
            time = calc(_edge)

        x.append(time)
        y.append(_edge.to_tl.temp)
        return array(x), array(y), min(x), max(x), min(y), max(y)


    @staticmethod
    def unpack(filename, parent):
        xmltree = etree.parse(filename)
        root = xmltree.getroot()
        solder_node = root[0]
        s = Solder(filename, solder_node.attrib["name"], solder_node.attrib["description"], parent)
        env_count = 0
        for temp_level in solder_node.findall("state"):
            tstr = temp_level.attrib["temperature"]
            is_env = False
            try:
                temp = int(tstr)
            except ValueError:
                if tstr == "$ENV":
                    temp = getTemperature()
                    is_env = True
                    env_count += 1
            s.add_temp_level(temp_level.attrib["name"], temp, is_env)

        set_colors(s.temp_levels)

        for edge in solder_node.findall("edge"):
            from_tl = s.get_temp_level_by_name(edge.attrib["from"])
            to_tl = s.get_temp_level_by_name(edge.attrib["to"])
            duration = None
            rate = None
            try:
                duration = float(edge.attrib["duration"])
            except ValueError:
                pass
            try:
                rate = float(edge.attrib["rate"])
            except ValueError:
                pass

            e = Edge(from_tl, to_tl, duration, rate)
            s.edges.append(e)
        return s

    def save(self):
        if self.changed:
            solder_node = etree.Element("solder_type", {"name" : self.name, "description" : self.description})
            for temp_level in self.temp_levels:
                temp = temp_level.is_env and "$ENV" or str(temp_level.temp)
                solder_node.append(etree.Element("state", {"name" : temp_level.name, "temperature" : temp}))
            for edge in self.edges:
                element = etree.Element("edge", {
                    "from" : edge.from_tl.name,
                    "to" : edge.to_tl.name,
                    "duration" : str(edge.duration),
                    "rate" : str(edge.rate)})
                solder_node.append(element)

            dirname = SolderListModel.dirname()
            root = etree.Element("xml")
            root.append(solder_node)
            if self.filename is None:
                self.filename = os.path.join(dirname, self.name + ".xml")
            etree.ElementTree(root).write(self.filename, "UTF-8", True)
            self.changed = False

    def setChanged(self):
        self.changed = True


class SolderListModel(QtCore.QAbstractListModel):
    def __init__(self, parent=None, *args):
        """ datain: a list where each item is a row
        """
        super(SolderListModel, self).__init__(parent, *args)
        self._load_solder_list()

    def rowCount(self, parent=QtCore.QModelIndex()):
        return len(self.solder_list)

    def _load_solder_list(self):
        dirname = self.dirname()
        dirlisting = filter(lambda x: os.path.splitext(x)[1] == ".xml", os.listdir(dirname))
        self.solder_list = []
        for p in dirlisting:
            self.solder_list.append(Solder.unpack(os.path.join(dirname, p), self))
        self.solder_list.sort(key=lambda x: x.name)
        self.reset()

    def headerData(self, col, orientation, role):
        if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
            return QtCore.QVariant("Solder Paste")
        return QtCore.QVariant()

    def data(self, index, role):
        if not index.isValid():
            return QtCore.QVariant()

        solder = self.solder_list[index.row()]
        if role == QtCore.Qt.DisplayRole:
            return QtCore.QVariant(solder.name)

        elif role == QtCore.Qt.DecorationRole and solder.changed:
            return QtGui.QIcon.fromTheme("document-save")

    def setData(self, index, variant, role):
        if index.isValid() and role == QtCore.Qt.EditRole:
            new_name = str(variant.toString())
            if new_name and new_name != "":
                self.solder_list[index.row()].name = new_name
                self.solder_list[index.row()].changed = True
                return True
        return False

    def flags(self, index):
        if not index.isValid():
            return 0
        return QtCore.Qt.ItemFlags(QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEditable)

    def create_solder(self):
        solder = Solder(None, datetime.now().strftime("%Y-%M-%D %H:%m:%s"), "")

        tl = solder.add_temp_level("environment temp", getTemperature(), True)
        tl.color = QtGui.QColor(0, 0, 0)
        self.solder_list.append(solder)
        self.reset()

    @staticmethod
    def dirname():
        return os.path.join(os.path.dirname(__file__), "solder_types")

    def check_name(self, name):
        for solder in self.solder_list:
            if name == solder.name:
                return False
        return True


class SolderWidget(QtGui.QWidget):
    def __init__(self, parent, readonly=False):
        super(SolderWidget, self).__init__(parent)
        self.solder_model = SolderListModel(self)
        self.solder_view = QtGui.QListView()
        self.solder_view.setModel(self.solder_model)
        self.solder_controls = AddRemoveWidget(self)

        layout = QtGui.QHBoxLayout(self)
        layout.addWidget(self.solder_view)
        layout.addWidget(self.solder_controls)

        self.connect(
            self.solder_controls.add_button,
            QtCore.SIGNAL("clicked()"),
            self.solder_model.create_solder)

        self.connect(
            self.solder_controls.remove_button,
            QtCore.SIGNAL("clicked()"),
            self.remove_solder)

        self.solder_view.setCurrentIndex(self.solder_model.index(0,0))

    def remove_solder(self):
        index = self.solder_view.currentIndex()
        solder = self.solder_model.solder_list[index.row()]
        try:
            os.remove(solder.filename)
        except OSError:
            pass
        del self.solder_model.solder_list[index.row()]
        self.solder_model.reset()
        new_index = self.solder_model.index(0)
        self.solder_view.setCurrentIndex(new_index)
        return new_index

    def save_solder(self, index):
        self.solder_model.solder_list[index.row()].save()
        self.solder_model.reset()
        new_index = self.solder_model.index(self.solder_model.solder_list.index(self.plotter.solder))
        self.solder_widget.solder_view.setCurrentIndex(new_index)