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.

def indent(elem, level=0):
    i = "\n" + level*"  "
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "  "
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for child in elem:
            indent(child, level+1)
        if not child.tail or not child.tail.strip():
            child.tail = i
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i


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):
        if not self.edges:
            return array([]), array([]), 0, 300, 0, 300
        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(self.edges[-1].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):
        self.changed = True
        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")
            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.solder_list = []
        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))
        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 is not None and new_name != "" and
                self.check_name(new_name)):
                solder = self.solder_list[index.row()]
                solder.name = new_name
                solder.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, "new %d" % len(self.solder_list), "")

        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):
    solder_selected = QtCore.pyqtSignal(QtCore.QModelIndex)

    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.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 create_solder(self):
        self.solder_model.create_solder()
        ix = self.solder_model.index(self.solder_model.rowCount(None) - 1, 0)
        self.solder_view.setCurrentIndex(ix)
        self.solder_selected.emit(ix)

    def remove_solder(self):
        index = self.solder_view.currentIndex()
        solder = self.solder_model.solder_list[index.row()]
        try:
            os.remove(solder.filename)
        except (OSError, TypeError):
            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)