datamodel.py

"""Implementation of the Datamodel for the parts list with natural sort."""

import re

import wx
import wx.dataview as dv


BOM_COL = 6
POS_COL = 7

class PartListDataModel(dv.PyDataViewModel):
    """Datamodel for use with the DataViewCtrl of the mainwindow."""

    def __init__(self):
        super().__init__()
        self.data = []

    @staticmethod
    def natural_sort_key(s):
        """Return a tuple that can be used for natural sorting."""
        return [
            int(text) if text.isdigit() else text.lower()
            for text in re.split("([0-9]+)", s)
        ]

    def GetColumnCount(self):
        return 10

    def GetColumnType(self, col):
        columntypes = (
            "string",
            "string",
            "string",
            "string",
            "string",
            "string",
            "wxDataViewIconText",
            "wxDataViewIconText",
            "string",
            "string",
        )
        return columntypes[col]

    def GetChildren(self, parent, children):
        if not parent:
            for row in self.data:
                children.append(self.ObjectToItem(row))
            return len(self.data)
        return 0

    def IsContainer(self, item):
        return not item

    def GetParent(self, item):
        return dv.NullDataViewItem

    def GetValue(self, item, col):
        row = self.ItemToObject(item)
        if col in [BOM_COL, POS_COL]:
            icon = row[col]
            return dv.DataViewIconText("", icon)
        return row[col]

    def SetValue(self, value, item, col):
        row = self.ItemToObject(item)
        if col in [BOM_COL, POS_COL]:
            return False
        row[col] = value
        return True

    def Compare(self, item1, item2, column, ascending):
        """Override to implement natural sorting."""
        val1 = self.GetValue(item1, column)
        val2 = self.GetValue(item2, column)

        key1 = self.natural_sort_key(val1)
        key2 = self.natural_sort_key(val2)

        if ascending:
            return (key1 > key2) - (key1 < key2)
        else:
            return (key2 > key1) - (key2 < key1)

    @staticmethod
    def get_bom_pos_icon(state):
        icons = [wx.ArtProvider.GetBitmap(wx.ART_INFORMATION, wx.ART_OTHER, wx.Size(16, 16)),
                 wx.ArtProvider.GetBitmap(wx.ART_WARNING, wx.ART_OTHER, wx.Size(16, 16))]
        return icons[int(state)]

    def AddEntry(self, item):
        """Method to add a new entry to the data model."""
        item[BOM_COL] = self.get_bom_pos_icon(item[BOM_COL])
        item[POS_COL] = self.get_bom_pos_icon(item[POS_COL])
        self.data.append(item)
        self.ItemAdded(dv.NullDataViewItem, self.ObjectToItem(item))

    def UpdateEntry(self, index, data):
        """Method to update an entry in the data model."""
        if index < 0 or index >= len(self.data):
            return
        item = self.data[index]
        # ToDo: Figure out why item = data does not work
        for i in range(0,len(data)):
            if i in [BOM_COL, POS_COL]:
                item[i] = self.get_bom_pos_icon(data[i])
            else:
                item[i] = data[i]
        self.ItemChanged(self.ObjectToItem(item))

    def RemoveEntry(self, item):
        """Method to remove an entry from the data model.""" 
        obj = self.ItemToObject(item)
        if obj in self.data:
            self.data.remove(obj)
            self.ItemDeleted(dv.NullDataViewItem, item)