[pygobject] Add Pythonic iterators and indexing to GVariant

[Martin Pitt <martinpitt src gnome org>]

commit b1a98083cdc50653e1d7bfb809bdf089f833df3d
Author: Martin Pitt <martin pitt ubuntu com>
Date:   Tue Jan 18 12:01:28 2011 +0100

    Add Pythonic iterators and indexing to GVariant
    
    Add the usual set of iterators and index accessors to GLib.Variant objects
    which are containers.
    
    Add corresponding test cases.

 gi/overrides/GLib.py    |   50 ++++++++++++++++++++++++++++++++++++++++++
 tests/test_overrides.py |   55 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+), 0 deletions(-)
---
diff --git a/gi/overrides/GLib.py b/gi/overrides/GLib.py
index 4cfdf52..38453de 100644
--- a/gi/overrides/GLib.py
+++ b/gi/overrides/GLib.py
@@ -195,6 +195,56 @@ class Variant(GLib.Variant):
 
         raise NotImplementedError, 'unsupported GVariant type ' + self.get_type_string()
 
+    #
+    # Pythonic iterators
+    #
+    def __len__(self):
+        if self.get_type_string().startswith('a') or self.get_type_string().startswith('('):
+            return self.n_children()
+        raise TypeError, 'GVariant type %s is not a container' % self.get_type_string()
+
+    def __getitem__(self, key):
+        # dict
+        if self.get_type_string().startswith('a{'):
+            try:
+                val = self.lookup_value(key, variant_type_from_string('*'))
+                if val is None:
+                    raise KeyError, key
+                return val.unpack()
+            except TypeError:
+                # lookup_value() only works for string keys, which is certainly
+                # the common case; we have to do painful iteration for other
+                # key types
+                for i in xrange(self.n_children()):
+                    v = self.get_child_value(i)
+                    if v.get_child_value(0).unpack() == key:
+                        return v.get_child_value(1).unpack()
+                raise KeyError, key
+
+        # array/tuple
+        if self.get_type_string().startswith('a') or self.get_type_string().startswith('('):
+            try:
+                key = int(key)
+            except ValueError, e:
+                raise TypeError, str(e)
+            if key < 0:
+                key = self.n_children() + key
+            if key < 0 or key >= self.n_children():
+                raise IndexError, 'list index out of range'
+            return self.get_child_value(key).unpack()
+
+        raise TypeError, 'GVariant type %s is not a container' % self.get_type_string()
+
+    def keys(self):
+        if not self.get_type_string().startswith('a{'):
+            return TypeError, 'GVariant type %s is not a dictionary' % self.get_type_string()
+
+        res = []
+        for i in xrange(self.n_children()):
+            v = self.get_child_value(i)
+            res.append(v.get_child_value(0).unpack())
+        return res
+
 @classmethod
 def new_tuple(cls, *elements):
     return variant_new_tuple(elements)
diff --git a/tests/test_overrides.py b/tests/test_overrides.py
index bbcd1fc..1655d20 100644
--- a/tests/test_overrides.py
+++ b/tests/test_overrides.py
@@ -69,6 +69,61 @@ class TestGLib(unittest.TestCase):
         res = GLib.Variant('a{si}', {'key1': 1, 'key2': 2}).unpack()
         self.assertEqual(res, {'key1': 1, 'key2': 2})
 
+    def test_gvariant_iteration(self):
+        # array index access
+        vb = GLib.VariantBuilder()
+        vb.init(gi._gi.variant_type_from_string('ai'))
+        vb.add_value(GLib.Variant.new_int32(-1))
+        vb.add_value(GLib.Variant.new_int32(3))
+        v = vb.end()
+
+        self.assertEqual(len(v), 2)
+        self.assertEqual(v[0], -1)
+        self.assertEqual(v[1], 3)
+        self.assertEqual(v[-1], 3)
+        self.assertEqual(v[-2], -1)
+        self.assertRaises(IndexError, v.__getitem__, 2)
+        self.assertRaises(IndexError, v.__getitem__, -3)
+        self.assertRaises(TypeError, v.__getitem__, 'a')
+
+        # array iteration
+        self.assertEqual([x for x in v], [-1, 3])
+        self.assertEqual(list(v), [-1, 3])
+
+        # tuple index access
+        v = GLib.Variant.new_tuple(GLib.Variant.new_int32(-1),
+                GLib.Variant.new_string('hello'))
+        self.assertEqual(len(v), 2)
+        self.assertEqual(v[0], -1)
+        self.assertEqual(v[1], 'hello')
+        self.assertEqual(v[-1], 'hello')
+        self.assertEqual(v[-2], -1)
+        self.assertRaises(IndexError, v.__getitem__, 2)
+        self.assertRaises(IndexError, v.__getitem__, -3)
+        self.assertRaises(TypeError, v.__getitem__, 'a')
+
+        # tuple iteration
+        self.assertEqual([x for x in v], [-1, 'hello'])
+        self.assertEqual(tuple(v), (-1, 'hello'))
+
+        # dictionary index access
+        vsi = GLib.Variant('a{si}', {'key1': 1, 'key2': 2})
+        vis = GLib.Variant('a{is}', {1: 'val1', 5: 'val2'})
+
+        self.assertEqual(len(vsi), 2)
+        self.assertEqual(vsi['key1'], 1)
+        self.assertEqual(vsi['key2'], 2)
+        self.assertRaises(KeyError, vsi.__getitem__, 'unknown')
+
+        self.assertEqual(len(vis), 2)
+        self.assertEqual(vis[1], 'val1')
+        self.assertEqual(vis[5], 'val2')
+        self.assertRaises(KeyError, vsi.__getitem__, 3)
+
+        # dictionary iteration
+        self.assertEqual(set(vsi.keys()), set(['key1', 'key2']))
+        self.assertEqual(set(vis.keys()), set([1, 5]))
+
 class TestPango(unittest.TestCase):
 
     def test_font_description(self):