[pygobject] Make TreeModel behave like in GTK-2.x

[John Palmieri <johnp src gnome org>]

commit a87e3ba64b54e6df0b5b96af47c34e3be790b58f
Author: Sebastian Pölsterl <sebp k-d-w org>
Date:   Thu Oct 7 19:37:53 2010 +0200

    Make TreeModel behave like in GTK-2.x
    
    Moved stuff from __getitem__ to get_iter.
    Added TreePath.__cmp__
    
    get_iter_from_string throws ValueError.
    iterchildren() does not return None.
    
    Adjusted tests to new TreeModel and added TestGtk.test_tree_model method
    
    Added support for negative row and column indices
    
    Use rich comparison methods instead of __cmp__
    
    Added TreeModel.__bool__/__nonzero__
    
    Raise Error if tree path string is empty
    
    https://bugzilla.gnome.org/show_bug.cgi?id=631547

 gi/overrides/Gtk.py     |  195 +++++++++++++++++++++++++++++++++++++++++++++++
 tests/test_overrides.py |  160 +++++++++++++++++++++++++++++++++++++--
 2 files changed, 349 insertions(+), 6 deletions(-)
---
diff --git a/gi/overrides/Gtk.py b/gi/overrides/Gtk.py
index 20c01a6..5988166 100644
--- a/gi/overrides/Gtk.py
+++ b/gi/overrides/Gtk.py
@@ -358,6 +358,89 @@ class TreeModel(Gtk.TreeModel):
     def __len__(self):
         return self.iter_n_children(None)
 
+    def __bool__(self):
+        return True
+
+    def __getitem__(self, key):
+        if isinstance(key, Gtk.TreeIter):
+            return TreeModelRow(self, key)
+        elif isinstance(key, int) and key < 0:
+            index = len(self) + key
+            if index < 0:
+                raise IndexError("row index is out of bounds: %d" % key)
+            try:
+                aiter = self.get_iter(index)
+            except ValueError:
+                raise IndexError("could not find tree path '%s'" % key)
+            return TreeModelRow(self, aiter)
+        else:
+            try:
+                aiter = self.get_iter(key)
+            except ValueError:
+                raise IndexError("could not find tree path '%s'" % key)
+            return TreeModelRow(self, aiter)
+
+    def __iter__(self):
+        return TreeModelRowIter(self, self.get_iter_first())
+
+    def get_iter(self, path):
+        if isinstance(path, Gtk.TreePath):
+            pass
+        elif isinstance(path, (int, str,)):
+            path = self._tree_path_from_string(str(path))
+        elif isinstance(path, tuple):
+            path_str = ":".join(str(val) for val in path)
+            path = self._tree_path_from_string(path_str)
+        else:
+            raise TypeError("tree path must be one of Gtk.TreeIter, Gtk.TreePath, \
+                int, str or tuple, not %s" % type(path).__name__)
+
+        success, aiter = super(TreeModel, self).get_iter(path)
+        if not success:
+            raise ValueError("invalid tree path '%s'" % path)
+        return aiter
+
+    def _tree_path_from_string(self, path):
+        if len(path) == 0:
+            raise TypeError("could not parse subscript '%s' as a tree path" % path)
+        try:
+            return TreePath.new_from_string(path)
+        except TypeError:
+            raise TypeError("could not parse subscript '%s' as a tree path" % path)
+
+    def get_iter_first(self):
+        success, aiter = super(TreeModel, self).get_iter_first()
+        if success:
+            return aiter
+
+    def get_iter_from_string(self, path_string):
+        success, aiter = super(TreeModel, self).get_iter_from_string(path_string)
+        if not success:
+            raise ValueError("invalid tree path '%s'" % path_string)
+        return aiter
+
+    def iter_next(self, aiter):
+        next_iter = aiter.copy()
+        success = super(TreeModel, self).iter_next(next_iter)
+        if success:
+            return next_iter
+
+    def iter_children(self, aiter):
+        success, child_iter = super(TreeModel, self).iter_children(aiter)
+        if success:
+            return child_iter
+
+    def iter_nth_child(self, parent, n):
+        success, child_iter = super(TreeModel, self).iter_nth_child(parent, n)
+        if success:
+            return child_iter
+
+    def iter_parent(self, aiter):
+        success, parent_iter = super(TreeModel, self).iter_parent(aiter)
+        if success:
+            return parent_iter
+
+TreeModel.__nonzero__ = TreeModel.__bool__
 TreeModel = override(TreeModel)
 __all__.append('TreeModel')
 
@@ -386,6 +469,118 @@ class ListStore(Gtk.ListStore, TreeModel):
 ListStore = override(ListStore)
 __all__.append('ListStore')
 
+class TreeModelRow(object):
+
+    def __init__(self, model, iter_or_path):
+        if not isinstance(model, Gtk.TreeModel):
+            raise TypeError("expected Gtk.TreeModel, %s found" % type(model).__name__)
+        self.model = model
+        if isinstance(iter_or_path, Gtk.TreePath):
+            self.iter = model.get_iter(iter_or_path)
+        elif isinstance(iter_or_path, Gtk.TreeIter):
+            self.iter = iter_or_path
+        else:
+            raise TypeError("expected Gtk.TreeIter or Gtk.TreePath, \
+                %s found" % type(iter_or_path).__name__)
+
+    @property
+    def path(self):
+        return self.model.get_path(self.iter)
+
+    @property
+    def next(self):
+        return self.get_next()
+
+    @property
+    def parent(self):
+        return self.get_parent()
+
+    def get_next(self):
+        next_iter = self.model.iter_next(self.iter)
+        if next_iter:
+            return TreeModelRow(self.model, next_iter)
+
+    def get_parent(self):
+        parent_iter = self.model.iter_parent(self.iter)
+        if parent_iter:
+            return TreeModelRow(self.model, parent_iter)
+
+    def __getitem__(self, key):
+        if isinstance(key, int):
+            if key >= self.model.get_n_columns():
+                raise IndexError("column index is out of bounds: %d" % key)
+            elif key < 0:
+                key = self._convert_negative_index(key)
+            return self.model.get_value(self.iter, key)
+        else:
+            raise TypeError("indices must be integers, not %s" % type(key).__name__)
+
+    def __setitem__(self, key, value):
+        if isinstance(key, int):
+            if key >= self.model.get_n_columns():
+                raise IndexError("column index is out of bounds: %d" % key)
+            elif key < 0:
+                key = self._convert_negative_index(key)
+            return self.model.set_value(self.iter, key, value)
+        else:
+            raise TypeError("indices must be integers, not %s" % type(key).__name__)
+
+    def _convert_negative_index(self, index):
+        new_index = self.model.get_n_columns() + index
+        if new_index < 0:
+            raise IndexError("column index is out of bounds: %d" % index)
+        return new_index
+
+    def iterchildren(self):
+        child_iter = self.model.iter_children(self.iter)
+        return TreeModelRowIter(self.model, child_iter)
+
+__all__.append('TreeModelRow')
+
+class TreeModelRowIter(object):
+
+    def __init__(self, model, aiter):
+        self.model = model
+        self.iter = aiter
+
+    def next(self):
+        if not self.iter:
+            raise StopIteration
+        row = TreeModelRow(self.model, self.iter)
+        self.iter = self.model.iter_next(self.iter)
+        return row
+
+    def __iter__(self):
+        return self
+
+__all__.append('TreeModelRowIter')
+
+class TreePath(Gtk.TreePath):
+
+    def __str__(self):
+        return self.to_string()
+
+    def __lt__(self, other):
+        return self.compare(other) < 0
+
+    def __le__(self, other):
+        return self.compare(other) <= 0
+
+    def __eq__(self, other):
+        return self.compare(other) == 0
+
+    def __ne__(self, other):
+        return self.compare(other) != 0
+
+    def __gt__(self, other):
+        return self.compare(other) > 0
+
+    def __ge__(self, other):
+        return self.compare(other) >= 0
+
+TreePath = override(TreePath)
+__all__.append('TreePath')
+
 class TreeStore(Gtk.TreeStore, TreeModel):
 
     def __init__(self, *column_types):
diff --git a/tests/test_overrides.py b/tests/test_overrides.py
index b1e3617..ebb2eff 100644
--- a/tests/test_overrides.py
+++ b/tests/test_overrides.py
@@ -226,14 +226,14 @@ class TestGtk(unittest.TestCase):
         parent = None
         i = 0
 
-        (has_children, treeiter) = tree_store.iter_children(parent)
-        while (has_children):
+        treeiter = tree_store.iter_children(parent)
+        while treeiter:
            i = tree_store.get_value(treeiter, 0)
            s = tree_store.get_value(treeiter, 1)
            obj = tree_store.get_value(treeiter, 2)
            obj.check(i, s)
            parent = treeiter
-           (has_children, treeiter) = tree_store.iter_children(parent)
+           treeiter = tree_store.iter_children(parent)
 
         self.assertEquals(i, 99)
 
@@ -248,17 +248,165 @@ class TestGtk(unittest.TestCase):
 
         # walk the list to see if the values were stored correctly
         i = 0
-        (has_more, treeiter) = list_store.get_iter_first()
+        treeiter = list_store.get_iter_first()
 
-        while has_more:
+        while treeiter:
             i = list_store.get_value(treeiter, 0)
             s = list_store.get_value(treeiter, 1)
             obj = list_store.get_value(treeiter, 2)
             obj.check(i, s)
-            has_more = list_store.iter_next(treeiter)
+            treeiter = list_store.iter_next(treeiter)
 
         self.assertEquals(i, 99)
 
+    def test_tree_model(self):
+        tree_store = Gtk.TreeStore(int, str)
+
+        self.assertTrue(tree_store)
+        self.assertEqual(len(tree_store), 0)
+        self.assertEqual(tree_store.get_iter_first(), None)
+
+        def get_by_index(row, col=None):
+            if col:
+                return tree_store[row][col]
+            else:
+                return tree_store[row]
+
+        self.assertRaises(TypeError, get_by_index, None)
+        self.assertRaises(TypeError, get_by_index, "")
+        self.assertRaises(TypeError, get_by_index, ())
+
+        self.assertRaises(IndexError, get_by_index, "0")
+        self.assertRaises(IndexError, get_by_index, 0)
+        self.assertRaises(IndexError, get_by_index, (0,))
+
+        self.assertRaises(ValueError, tree_store.get_iter, "0")
+        self.assertRaises(ValueError, tree_store.get_iter, 0)
+        self.assertRaises(ValueError, tree_store.get_iter, (0,))
+
+        self.assertRaises(ValueError, tree_store.get_iter_from_string, "0")
+
+        for row in tree_store:
+            self.fail("Should not be reached")
+
+        for i in range(100):
+            label = 'this is row #%d' % i
+            parent = tree_store.append(None, (i, label,))
+            self.assertNotEquals(parent, None)
+            for j in range(20):
+                label = 'this is child #%d of node #%d' % (j, i)
+                child = tree_store.append(parent, (j, label,))
+                self.assertNotEqual(child, None)
+
+        self.assertTrue(tree_store)
+        self.assertEqual(len(tree_store), 100)
+
+        for i,row in enumerate(tree_store):
+            self.assertEqual(row.model, tree_store)
+            self.assertEqual(row.parent, None)
+
+            self.assertEqual(tree_store[i].path, row.path)
+            self.assertEqual(tree_store[str(i)].path, row.path)
+            self.assertEqual(tree_store[(i,)].path, row.path)
+
+            self.assertEqual(tree_store[i][0], i)
+            self.assertEqual(tree_store[i][1], "this is row #%d" % i)
+
+            aiter = tree_store.get_iter(i)
+            self.assertEqual(tree_store.get_path(aiter), row.path)
+
+            aiter = tree_store.get_iter(str(i))
+            self.assertEqual(tree_store.get_path(aiter), row.path)
+
+            aiter = tree_store.get_iter((i,))
+            self.assertEqual(tree_store.get_path(aiter), row.path)
+
+            self.assertEqual(tree_store.iter_parent(aiter), row.parent)
+
+            next = tree_store.iter_next(aiter)
+            if i < len(tree_store) - 1:
+                self.assertEqual(tree_store.get_path(next), row.next.path)
+            else:
+                self.assertEqual(next, None)
+
+            self.assertEqual(tree_store.iter_n_children(row.iter), 20)
+
+            child = tree_store.iter_children(row.iter)
+            for j,childrow in enumerate(row.iterchildren()):
+                child_path = tree_store.get_path(child)
+                self.assertEqual(childrow.path, child_path)
+                self.assertEqual(childrow.parent.path, row.path)
+                self.assertEqual(childrow.path, tree_store[child].path)
+                self.assertEqual(childrow.path, tree_store[child_path].path)
+
+                self.assertEqual(childrow[0], tree_store[child][0])
+                self.assertEqual(childrow[0], j)
+                self.assertEqual(childrow[1], tree_store[child][1])
+                self.assertEqual(childrow[1], 'this is child #%d of node #%d' % (j, i))
+
+                self.assertRaises(IndexError, get_by_index, child, 2)
+
+                tree_store[child][1] = 'this was child #%d of node #%d' % (j, i)
+                self.assertEqual(childrow[1], 'this was child #%d of node #%d' % (j, i))
+
+                nth_child = tree_store.iter_nth_child(row.iter, j)
+                self.assertEqual(childrow.path, tree_store.get_path(nth_child))
+
+                childrow2 = tree_store["%d:%d" % (i, j)]
+                self.assertEqual(childrow.path, childrow2.path)
+
+                childrow2 = tree_store[(i, j,)]
+                self.assertEqual(childrow.path, childrow2.path)
+
+                child = tree_store.iter_next(child)
+                if j < 19:
+                    self.assertEqual(childrow.next.path, tree_store.get_path(child))
+                else:
+                    self.assertEqual(child, childrow.next)
+                    self.assertEqual(child, None)
+
+            self.assertEqual(j, 19)
+
+        self.assertEqual(i, 99)
+
+        # negative indices
+        for i in range(-1,-100,-1):
+            i_real = i + 100
+            self.assertEqual(tree_store[i][0], i_real)
+
+            row = tree_store[i]
+            for j in range(-1, -20, -1):
+                j_real = j + 20
+                path = (i_real, j_real,)
+
+                self.assertEqual(tree_store[path][-2], j_real)
+
+                label = 'this was child #%d of node #%d' % (j_real, i_real)
+                self.assertEqual(tree_store[path][-1], label)
+
+                new_label = 'this still is child #%d of node #%d' % (j_real, i_real)
+                tree_store[path][-1] = new_label
+                self.assertEqual(tree_store[path][-1], new_label)
+
+                self.assertRaises(IndexError, get_by_index, path, -3)
+
+        self.assertRaises(IndexError, get_by_index, -101)
+
+        last_row = tree_store[99]
+        self.assertNotEqual(last_row, None)
+
+        for i,childrow in enumerate(last_row.iterchildren()):
+            if i < 19:
+                self.assertTrue(tree_store.remove(childrow.iter))
+            else:
+                self.assertFalse(tree_store.remove(childrow.iter))
+
+        self.assertEqual(i, 19)
+
+        self.assertEqual(tree_store.iter_n_children(last_row.iter), 0)
+        for childrow in last_row.iterchildren():
+            self.fail("Should not be reached")
+
     def test_tree_view_column(self):
         cell = Gtk.CellRendererText()
         column = Gtk.TreeViewColumn(title='This is just a test',