Added support for __setitem__ to TreeModel and support for slices to TreeModelRow

https://bugzilla.gnome.org/show_bug.cgi?id=656891

2 files changed, 81 insertions, 2 deletions

diff --git a/gi/overrides/Gtk.py b/gi/overrides/Gtk.py
index 2aea6ad8..97c8b40b 100644
--- a/gi/overrides/Gtk.py
+++ b/gi/overrides/Gtk.py
@@ -734,6 +734,10 @@ class TreeModel(Gtk.TreeModel):
                 raise IndexError("could not find tree path '%s'" % key)
             return TreeModelRow(self, aiter)
 
+    def __setitem__(self, key, value):
+        row = self[key]
+        self.set_row(row.iter, value)
+
     def __iter__(self):
         return TreeModelRowIter(self, self.get_iter_first())
 
@@ -781,6 +785,8 @@ class TreeModel(Gtk.TreeModel):
     def set_row(self, treeiter, row):
         # TODO: Accept a dictionary for row
         # model.append(None,{COLUMN_ICON: icon, COLUMN_NAME: name})
+        if isinstance(row, str):
+            raise TypeError('Expected a list or tuple, but got str')
 
         n_columns = self.get_n_columns()
         if len(row) != n_columns:
@@ -1033,6 +1039,12 @@ class TreeModelRow(object):
             elif key < 0:
                 key = self._convert_negative_index(key)
             return self.model.get_value(self.iter, key)
+        elif isinstance(key, slice):
+            start, stop, step = key.indices(self.model.get_n_columns())
+            alist = []
+            for i in range(start, stop, step):
+                alist.append(self.model.get_value(self.iter, i))
+            return alist
         else:
             raise TypeError("indices must be integers, not %s" % type(key).__name__)
 
@@ -1042,9 +1054,19 @@ class TreeModelRow(object):
                 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)
+            self.model.set_value(self.iter, key, value)
+        elif isinstance(key, slice):
+            start, stop, step = key.indices(self.model.get_n_columns())
+            indexList = range(start, stop, step)
+            if len(indexList) != len(value):
+                raise ValueError(
+                    "attempt to assign sequence of size %d to slice of size %d"
+                        % (len(value), len(indexList)))
+
+            for i,v in enumerate(indexList):
+                self.model.set_value(self.iter, v, value[i])
         else:
-            raise TypeError("indices must be integers, not %s" % type(key).__name__)
+            raise TypeError("index must be an integer or slice, not %s" % type(key).__name__)
 
     def _convert_negative_index(self, index):
         new_index = self.model.get_n_columns() + index
diff --git a/tests/test_overrides.py b/tests/test_overrides.py
index 427f4d12..2f251210 100644
--- a/tests/test_overrides.py
+++ b/tests/test_overrides.py
@@ -1246,6 +1246,63 @@ class TestGtk(unittest.TestCase):
         self.assertRaises(ValueError, tree_store.get, aiter, 1, 100)
         self.assertEqual(tree_store.get(aiter, 0, 1), (10, 'this is row #10'))
 
+    def test_tree_model_edit(self):
+        model = Gtk.ListStore(int, str, float)
+        model.append([1, "one", -0.1])
+        model.append([2, "two", -0.2])
+
+        def set_row(value):
+            model[1] = value
+
+        self.assertRaises(TypeError, set_row, 3)
+        self.assertRaises(TypeError, set_row, "three")
+        self.assertRaises(ValueError, set_row, [])
+        self.assertRaises(ValueError, set_row, [3, "three"])
+
+        model[0] = (3, "three", -0.3)
+
+    def test_tree_row_slice(self):
+        model = Gtk.ListStore(int, str, float)
+        model.append([1, "one", -0.1])
+
+        self.assertEqual([1, "one", -0.1], model[0][:])
+        self.assertEqual([1, "one"], model[0][:2])
+        self.assertEqual(["one", -0.1], model[0][1:])
+        self.assertEqual(["one"], model[0][1:-1])
+        self.assertEqual([1], model[0][:-2])
+        self.assertEqual([], model[0][5:])
+        self.assertEqual([1, -0.1], model[0][0:3:2])
+
+        model[0][:] = (2, "two", -0.2)
+        self.assertEqual([2, "two", -0.2], model[0][:])
+
+        model[0][:2] = (3, "three")
+        self.assertEqual([3, "three", -0.2], model[0][:])
+
+        model[0][1:] = ("four", -0.4)
+        self.assertEqual([3, "four", -0.4], model[0][:])
+
+        model[0][1:-1] = ("five",)
+        self.assertEqual([3, "five", -0.4], model[0][:])
+
+        model[0][0:3:2] = (6, -0.6)
+        self.assertEqual([6, "five", -0.6], model[0][:])
+
+        def set_row1():
+            model[0][5:] = ("doesn't", "matter",)
+
+        self.assertRaises(ValueError, set_row1)
+
+        def set_row2():
+            model[0][:1] = (0, "zero", 0)
+
+        self.assertRaises(ValueError, set_row2)
+
+        def set_row3():
+            model[0][:2] = ("0", 0)
+
+        self.assertRaises(ValueError, set_row3)
+
     def test_tree_view_column(self):
         cell = Gtk.CellRendererText()
         column = Gtk.TreeViewColumn(title='This is just a test',