1 files changed, 37 insertions, 3 deletions
diff --git a/src/gui/tasks/table/model.py b/src/gui/tasks/table/model.py
index cb25b22..5013569 100644
--- a/src/gui/tasks/table/model.py
+++ b/src/gui/tasks/table/model.py
@@ -1,16 +1,20 @@
-from PyQt5 import QtCore, QtWidgets
+from PyQt5 import QtCore, QtWidgets, QtGui
 from PyQt5.QtCore import Qt
 
 from model.task import Task
+from model.difficulty import Difficulty
+from model.priority import Priority
+from model import difficulty, priority
 import time
 import math
 import util.array
 import util.range
 import gui.tasks.duration
+import gui.color
 
-columns = 4
+columns = 6
 
-headers = ['Age', 'Name', 'Duration', 'Tag']
+headers = ['Age', 'Name', 'Duration', 'Difficulty', 'Priority', 'Tag']
 
 default_sort = (0, Qt.AscendingOrder)
 
@@ -28,6 +32,8 @@ class TableModel(QtCore.QAbstractTableModel):
             return QtCore.QVariant()
 
     def data(self, index, role):
+        task = self._tasks[index.row()]
+
         if role == Qt.DisplayRole:
             task = self._tasks[index.row()]
             if index.column() == 0:
@@ -37,7 +43,15 @@ class TableModel(QtCore.QAbstractTableModel):
             elif index.column() == 2:
                 return gui.tasks.duration.format(task.duration)
             elif index.column() == 3:
+                return difficulty.format(task.difficulty)
+            elif index.column() == 4:
+                return priority.format(task.priority)
+            elif index.column() == 5:
                 return task.tag
+        elif role == Qt.ForegroundRole and index.column() == 3:
+            return QtGui.QBrush(difficulty_color(task.difficulty))
+        elif role == Qt.ForegroundRole and index.column() == 4:
+            return QtGui.QBrush(priority_color(task.priority))
         else:
             return QtCore.QVariant()
 
@@ -112,6 +126,10 @@ def sort_key(task: Task, row: int):
     elif row == 2:
         return task.duration
     elif row == 3:
+        return task.difficulty
+    elif row == 4:
+        return task.priority
+    elif row == 5:
         return task.tag.lower()
 
 def is_reversed(row: int, order: Qt.SortOrder) -> bool:
@@ -119,3 +137,19 @@ def is_reversed(row: int, order: Qt.SortOrder) -> bool:
         return order == Qt.AscendingOrder
     else:
         return order == Qt.DescendingOrder
+
+def difficulty_color(d: Difficulty) -> QtGui.QColor:
+    if d == Difficulty.LOW:
+        return gui.color.low_difficulty
+    elif d == Difficulty.MEDIUM:
+        return gui.color.medium_difficulty
+    elif d == Difficulty.HIGH:
+        return gui.color.high_difficulty
+
+def priority_color(p: Priority) -> QtGui.QColor:
+    if p == Priority.LOW:
+        return gui.color.low_priority
+    elif p == Priority.MEDIUM:
+        return gui.color.medium_priority
+    elif p == Priority.HIGH:
+        return gui.color.high_priority