[pypy-commit] pypy reflex-support: pythonized indexing and slicing for std::vector-like objects

[wlav]

Author: Wim Lavrijsen <WLavrijsen at lbl.gov>
Branch: reflex-support
Changeset: r53735:897bb52ddf19
Date: 2012-03-16 02:26 -0700
http://bitbucket.org/pypy/pypy/changeset/897bb52ddf19/

Log:	pythonized indexing and slicing for std::vector-like objects

diff --git a/pypy/module/cppyy/pythonify.py b/pypy/module/cppyy/pythonify.py
--- a/pypy/module/cppyy/pythonify.py
+++ b/pypy/module/cppyy/pythonify.py
@@ -280,6 +280,26 @@
     return scope
 
 
+# pythonization by decoration (move to their own file?)
+import types
+
+def python_style_getitem(self, idx):
+    # python-style indexing: check for size and allow indexing from the back
+    sz = len(self)
+    if idx < 0: idx = sz + idx
+    if idx < sz:
+        return self._getitem__unchecked(idx)
+    raise IndexError('index out of range: %d requested for %s of size %d' % (idx, str(self), sz))
+
+def python_style_sliceable_getitem(self, slice_or_idx):
+    if type(slice_or_idx) == types.SliceType:
+        nseq = self.__class__()
+        nseq += [python_style_getitem(self, i) \
+                    for i in range(*slice_or_idx.indices(len(self)))]
+        return nseq
+    else:
+        return python_style_getitem(self, slice_or_idx)
+
 _pythonizations = {}
 def _pythonize(pyclass):
 
@@ -290,20 +310,39 @@
 
     # map size -> __len__ (generally true for STL)
     if hasattr(pyclass, 'size') and \
-            not hasattr(pyclass,'__len__') and callable(pyclass.size):
+            not hasattr(pyclass, '__len__') and callable(pyclass.size):
         pyclass.__len__ = pyclass.size
 
+    # map push_back -> __iadd__ (generally true for STL)
+    if hasattr(pyclass, 'push_back') and not hasattr(pyclass, '__iadd__'):
+        def __iadd__(self, ll):
+            [self.push_back(x) for x in ll]
+            return self
+        pyclass.__iadd__ = __iadd__
+
     # map begin()/end() protocol to iter protocol
     if hasattr(pyclass, 'begin') and hasattr(pyclass, 'end'):
-        def __iter__(self):
-            iter = self.begin()
+        try:
             # TODO: make gnu-independent
-            while gbl.__gnu_cxx.__ne__(iter, self.end()):
-                yield iter.__deref__()
-                iter.__preinc__()
-            iter.destruct()
-            raise StopIteration
-        pyclass.__iter__ = __iter__
+            ne = gbl.__gnu_cxx.__ne__
+            def __iter__(self):
+                iter = self.begin()
+                while gbl.__gnu_cxx.__ne__(iter, self.end()):
+                    yield iter.__deref__()
+                    iter.__preinc__()
+                iter.destruct()
+                raise StopIteration
+            pyclass.__iter__ = __iter__
+        except AttributeError:
+            pass
+
+    # combine __getitem__ and __len__ to make a pythonized __getitem__
+    if hasattr(pyclass, '__getitem__') and hasattr(pyclass, '__len__'):
+        pyclass._getitem__unchecked = pyclass.__getitem__
+        if hasattr(pyclass, '__setitem__') and hasattr(pyclass, '__iadd__'):
+            pyclass.__getitem__ = python_style_sliceable_getitem
+        else:
+            pyclass.__getitem__ = python_style_getitem
 
     # string comparisons (note: CINT backend requires the simple name 'string')
     if pyclass.__name__ == 'std::basic_string<char>' or pyclass.__name__ == 'string':
diff --git a/pypy/module/cppyy/test/test_stltypes.py b/pypy/module/cppyy/test/test_stltypes.py
--- a/pypy/module/cppyy/test/test_stltypes.py
+++ b/pypy/module/cppyy/test/test_stltypes.py
@@ -131,6 +131,65 @@
 
         v.destruct()
 
+    def test05_push_back_iterables_with_iadd(self):
+        """Test usage of += of iterable on push_back-able container"""
+
+        import cppyy
+
+        v = cppyy.gbl.std.vector(int)()
+
+        v += [1, 2, 3]
+        assert len(v) == 3
+        assert v[0] == 1
+        assert v[1] == 2
+        assert v[2] == 3
+
+        v += (4, 5, 6)
+        assert len(v) == 6
+        assert v[3] == 4
+        assert v[4] == 5
+        assert v[5] == 6
+
+        raises(TypeError, v.__iadd__, (7, '8'))  # string shouldn't pass
+        assert len(v) == 7   # TODO: decide whether this should roll-back
+
+        v2 = cppyy.gbl.std.vector(int)()
+        v2 += [8, 9]
+        assert len(v2) == 2
+        assert v2[0] == 8
+        assert v2[1] == 9
+
+        v += v2
+        assert len(v) == 9
+        assert v[6] == 7
+        assert v[7] == 8
+        assert v[8] == 9
+
+    def test06_vector_indexing(self):
+        """Test python-style indexing to an std::vector<int>"""
+
+        import cppyy
+
+        v = cppyy.gbl.std.vector(int)()
+
+        for i in range(self.N):
+            v.push_back(i)
+
+        raises(IndexError, 'v[self.N]')
+        raises(IndexError, 'v[self.N+1]')
+
+        assert v[-1] == self.N-1
+        assert v[-2] == self.N-2
+
+        assert len(v[0:0]) == 0
+        assert v[1:2][0] == v[1]
+
+        v2 = v[2:-1]
+        assert len(v2) == self.N-3     # 2 off from start, 1 from end
+        assert v2[0] == v[2]
+        assert v2[-1] == v[-2]
+        assert v2[self.N-4] == v[-2]
+
 
 class AppTestSTLSTRING:
     def setup_class(cls):