Added background colors to the model and display them

1 files changed, 20 insertions, 6 deletions

diff --git a/model.py b/model.py
index 75ad6a8..2084773 100644
--- a/model.py
+++ b/model.py
@@ -24,7 +24,7 @@ from sugar import dispatch
 
 class MindMapModel(gtk.GenericTreeModel):
 
-    _COLUMN_TYPES = (str, str, long, long)
+    _COLUMN_TYPES = (str, str, long, long, str)
 
     def __init__(self):
         gobject.GObject.__init__(self)
@@ -65,6 +65,7 @@ class MindMapModel(gtk.GenericTreeModel):
             thought_dict['name'] = thought.name
             thought_dict['x'] = thought.x
             thought_dict['y'] = thought.y
+            thought_dict['color'] = thought.color
             thoughts.append(thought_dict)
 
         return cjson.encode({'thoughts': thoughts})
@@ -73,9 +74,10 @@ class MindMapModel(gtk.GenericTreeModel):
         thoughts = cjson.decode(data)['thoughts']
         for thought_dict in thoughts:
             thought = Thought(thought_dict['id'])
-            thought.name = thought_dict['name']
-            thought.x = thought_dict['x']
-            thought.y = thought_dict['y']
+            thought.name = thought_dict.get('name', None)
+            thought.x = thought_dict.get('x', None)
+            thought.y = thought_dict.get('y', None)
+            thought.color = thought_dict.get('color', None)
             self._add_thought(thought)
 
     # gtk.GenericTreeModel methods
@@ -151,18 +153,19 @@ class MindMapModel(gtk.GenericTreeModel):
 
 class Thought(object):
 
-    def __init__(self, thought_id, name=None, x=0, y=0):
+    def __init__(self, thought_id, name=None, x=0, y=0, color=None):
         if thought_id is None:
             raise ValueError('thought_id cannot be None')
         self._id = thought_id
         self._name = name
         self._x = x
         self._y = y
+        self._color = color
 
         self.changed = dispatch.Signal()
 
     def get_tuple(self):
-        return (self.id, self.name, self.x, self.y)
+        return (self.id, self.name, self.x, self.y, self.color)
 
     def set_id(self, new_id):
         if self._id == new_id:
@@ -208,3 +211,14 @@ class Thought(object):
 
     y = property(get_y, set_y)
 
+    def set_color(self, new_color):
+        if self._color == new_color:
+            return
+        self._color = new_color
+        self.changed.send(self)
+
+    def get_color(self):
+        return self._color
+
+    color = property(get_color, set_color)
+