Add links to the model

1 files changed, 115 insertions, 36 deletions

diff --git a/model.py b/model.py
index 2084773..cefa1c7 100644
--- a/model.py
+++ b/model.py
@@ -15,6 +15,7 @@
 # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 
 import logging
+from gettext import gettext as _
 
 import gobject
 import gtk
@@ -25,33 +26,46 @@ from sugar import dispatch
 class MindMapModel(gtk.GenericTreeModel):
 
     _COLUMN_TYPES = (str, str, long, long, str)
+    
+    THOUGHTS = 0
+    LINKS = 1
 
     def __init__(self):
         gobject.GObject.__init__(self)
 
-        self._next_thought_id = 0
+        self._next_element_id = 2
         self._thoughts_by_id = {}
         self._thoughts = []
+        self._links_by_id = {}
+        self._links = []
 
     def create_new_thought(self):
-        thought = Thought(self._next_thought_id)
+        thought = Thought(self._next_element_id)
+        self._next_element_id += 1
         self._add_thought(thought)
 
     def _add_thought(self, thought):
         thought.changed.connect(self.__thought_changed_cb)
         self._thoughts.append(thought)
         self._thoughts_by_id[thought.id] = thought
-        self._next_thought_id += 1
 
-        path = (thought.id,)
+        path = (self.THOUGHTS, len(self._thoughts) - 1)
         self.row_inserted(path, self.get_iter(path))
 
-    def get_thought(self, thought_id):
-        return self._thoughts_by_id[int(thought_id)]
+    def get_thought(self, path):
+        if isinstance(path, basestring):
+            path = path.split(':')
+            for i in range(len(path)):
+                path[i] = int(path[i])
+
+        if path[0] != self.THOUGHTS or path[1] >= len(path):
+            raise ValueError('Invalid path %r' % (path,))
+
+        return self._thoughts[path[1]]
 
     def __thought_changed_cb(self, **kwargs):
         thought = kwargs['sender']
-        path = (thought.id,)
+        path = (self.THOUGHTS, self._thoughts.index(thought))
         self.row_changed(path, self.get_iter(path))
 
     def get_thoughts(self):
@@ -74,6 +88,8 @@ class MindMapModel(gtk.GenericTreeModel):
         thoughts = cjson.decode(data)['thoughts']
         for thought_dict in thoughts:
             thought = Thought(thought_dict['id'])
+            self._next_element_id = max(self._next_element_id + 1,
+                                        int(thought_dict['id']) + 1)
             thought.name = thought_dict.get('name', None)
             thought.x = thought_dict.get('x', None)
             thought.y = thought_dict.get('y', None)
@@ -81,7 +97,10 @@ class MindMapModel(gtk.GenericTreeModel):
             self._add_thought(thought)
 
     # gtk.GenericTreeModel methods
-    # paths and rowrefs are the thought ids
+    # paths are 0 for the THOUGHTS category, 1 for the LINKS category and the
+    # position in their list for elements
+    # rowrefs are 0 for the THOUGHTS category, 1 for the LINKS category and the
+    # id for elements
     def on_get_flags(self):
         return gtk.TREE_MODEL_ITERS_PERSIST
 
@@ -92,64 +111,124 @@ class MindMapModel(gtk.GenericTreeModel):
         return self._COLUMN_TYPES[n]
 
     def on_get_iter(self, path):
-        logging.debug('on_get_iter %r %r' % (type(path), path))
-        if path[0] in self._thoughts_by_id:
-            return path[0]
+        #logging.debug('on_get_iter %r %r' % (type(path), path))
+
+        if len(path) == 1 and path[0] == self.THOUGHTS:
+            return self.THOUGHTS
+        elif len(path) == 1 and path[0] == self.LINKS:
+            return self.LINKS
+        elif path[0] == self.THOUGHTS:
+            element_list = self._thoughts
+        elif path[0] == self.LINKS:
+            element_list = self._links
         else:
             return None
 
+        if path[1] < len(element_list):
+            return element_list[path[1]].id
+        else:
+            return None            
+
     def on_get_path(self, rowref):
         logging.debug('on_get_path %r' % rowref)
-        return (rowref,)
+        if rowref == self.THOUGHTS:
+            return (self.THOUGHTS,)
+        elif rowref == self.LINKS:
+            return (self.LINKS,)
+        elif rowref in self._thoughts_by_id:
+            index = self._thoughts.index(self._thoughts_by_id[rowref])
+            return (self.THOUGHTS, index,)
+        elif rowref in self._links_by_id:
+            index = self._links.index(self._links_by_id[rowref])
+            return (self.LINKS, index,)
+        else:
+            return None
 
     def on_get_value(self, rowref, column):
         #logging.debug('on_get_value %r %r' % (rowref, column))
-        if rowref >= len(self._thoughts):
+        if rowref == self.THOUGHTS and column == 1:
+            return _('Thoughts')
+        elif rowref == self.LINKS and column == 1:
+            return _('Links')
+        elif rowref in self._thoughts_by_id:
+            element = self._thoughts_by_id[rowref]
+        elif rowref in self._links_by_id:
+            element = self._links_by_id[rowref]
+        else:
             return None
-        thought = self._thoughts[rowref]
-        value_tuple = thought.get_tuple()
+
+        value_tuple = element.get_tuple()
         return value_tuple[column]
 
     def on_iter_next(self, rowref):
         #logging.debug('on_iter_next %r' % rowref)
-        if rowref not in self._thoughts_by_id:
+        if rowref == self.THOUGHTS:
+            return self.LINKS
+        if rowref in self._thoughts_by_id:
+            element = self._thoughts_by_id[rowref]
+            element_list = self._thoughts
+        elif rowref in self._links_by_id:
+            element = self._links_by_id[rowref]
+            element_list = self._links
+        else:
             return None
-        thought = self._thoughts_by_id[rowref]
         
-        index = self._thoughts.index(thought)
-        if index + 1 >= len(self._thoughts):
+        index = element_list.index(element)
+        if index + 1 >= len(element_list):
             return None
-        next_thought = self._thoughts[index + 1]
-        #logging.debug('on_iter_next returning %r' % next_thought.id)
-        return next_thought.id
+        next_element = element_list[index + 1]
+        return next_element.id
 
     def on_iter_children(self, rowref):
-        logging.debug('on_iter_children %r' % rowref)
-        if rowref:
-            return None
-        if self._thoughts:
+        #logging.debug('on_iter_children %r' % rowref)
+        if rowref is None:
+            return self.THOUGHTS
+        elif rowref == self.THOUGHTS and self._thoughts:
             return self._thoughts[0].id
+        elif rowref == self.LINKS and self._links:
+            return self._links[0].id
         else:
             return None
 
     def on_iter_has_child(self, rowref):
-        return False
+        #logging.debug('on_iter_has_child %r' % rowref)
+        if rowref == self.THOUGHTS and self._thoughts:
+            return True
+        elif rowref == self.LINKS and self._links:
+            return True
+        else:
+            return False
 
     def on_iter_n_children(self, rowref):
-        logging.debug('on_iter_n_children %r' % rowref)
-        if rowref:
+        #logging.debug('on_iter_n_children %r' % rowref)
+        if rowref == self.THOUGHTS:
+            return len(self._thoughts)
+        elif rowref == self.LINKS:
+            return len(self._links)
+        else:
             return 0
-        return len(self._thoughts)
 
     def on_iter_nth_child(self, rowref, n):
-        logging.debug('on_iter_nth_child %r %r' % (rowref, n))
-        if rowref is not None:
-            return None
-        if n < len(self._thoughts):
+        #logging.debug('on_iter_nth_child %r %r' % (rowref, n))
+        if rowref is None and n == self.THOUGHTS:
+            return self.THOUGHTS
+        elif rowref is None and n == self.LINKS:
+            return self.LINKS
+        elif rowref == self.THOUGHTS and n < len(self._thoughts):
             return self._thoughts[n].id
+        elif rowref == self.LINKS and n < len(self._links):
+            return self._links[n].id
+        else:
+            return None
 
-    def on_iter_parent(child):
-        return None
+    def on_iter_parent(self, rowref):
+        #logging.debug('on_iter_parent %r' % rowref)
+        if rowref in self._thoughts_by_id:
+            return self.THOUGHTS
+        elif rowref in self._links_by_id:
+            return self.LINKS
+        else:
+            return None
 
 class Thought(object):