Raise IndexError in PyQGIS when calling geometryN on a collection

nyalldawson · nyalldawson · commit f5a6aef55f3f · 2018-11-20T05:05:37.000+11:00
with an invalid geometry index

And add len operator to QgsGeometryCollection
diff --git a/python/core/auto_generated/geometry/qgsgeometrycollection.sip.in b/python/core/auto_generated/geometry/qgsgeometrycollection.sip.in
@@ -42,12 +42,40 @@ Returns the number of geometries within the collection.
 %End
 
 
-    QgsAbstractGeometry *geometryN( int n );
+    int __len__() const;
+%Docstring
+Returns the number of geometries within the collection.
+%End
+%MethodCode
+    sipRes = sipCpp->numGeometries();
+%End
+
+    //! Ensures that bool(obj) returns true (otherwise __len__() would be used)
+    int __bool__() const;
+%MethodCode
+    sipRes = true;
+%End
+
+
+
+    SIP_PYOBJECT geometryN( int n ) const;
 %Docstring
 Returns a geometry from within the collection.
 
 :param n: index of geometry to return
 %End
+%MethodCode
+    if ( a0 < 0 || a0 >= sipCpp->numGeometries() )
+    {
+      PyErr_SetString( PyExc_IndexError, QByteArray::number( a0 ) );
+      sipIsErr = 1;
+    }
+    else
+    {
+      return sipConvertFromType( sipCpp->geometryN( a0 ), sipType_QgsAbstractGeometry, NULL );
+    }
+%End
+
 
     virtual bool isEmpty() const;
 
diff --git a/src/core/geometry/qgsgeometrycollection.h b/src/core/geometry/qgsgeometrycollection.h
@@ -53,6 +53,24 @@ class CORE_EXPORT QgsGeometryCollection: public QgsAbstractGeometry
       return mGeometries.size();
     }
 
+#ifdef SIP_RUN
+
+    /**
+     * Returns the number of geometries within the collection.
+     */
+    int __len__() const;
+    % MethodCode
+    sipRes = sipCpp->numGeometries();
+    % End
+
+    //! Ensures that bool(obj) returns true (otherwise __len__() would be used)
+    int __bool__() const;
+    % MethodCode
+    sipRes = true;
+    % End
+#endif
+
+
     /**
      * Returns a const reference to a geometry from within the collection.
      * \param n index of geometry to return
@@ -67,7 +85,23 @@ class CORE_EXPORT QgsGeometryCollection: public QgsAbstractGeometry
      * Returns a geometry from within the collection.
      * \param n index of geometry to return
      */
+#ifndef SIP_RUN
     QgsAbstractGeometry *geometryN( int n );
+#else
+    SIP_PYOBJECT geometryN( int n ) const;
+    % MethodCode
+    if ( a0 < 0 || a0 >= sipCpp->numGeometries() )
+    {
+      PyErr_SetString( PyExc_IndexError, QByteArray::number( a0 ) );
+      sipIsErr = 1;
+    }
+    else
+    {
+      return sipConvertFromType( sipCpp->geometryN( a0 ), sipType_QgsAbstractGeometry, NULL );
+    }
+    % End
+#endif
+
 
     //methods inherited from QgsAbstractGeometry
     bool isEmpty() const override;
diff --git a/tests/src/python/test_qgsgeometry.py b/tests/src/python/test_qgsgeometry.py
@@ -242,6 +242,20 @@ def testReferenceGeometry(self):
                 result = geom.constGet().perimeter()
                 self.assertAlmostEqual(result, exp, 5, "Perimeter {}: mismatch Expected:\n{}\nGot:\n{}\n".format(i + 1, exp, result))
 
+    def testCollection(self):
+        g = QgsGeometry.fromWkt('MultiLineString()')
+        self.assertEqual(len(g.get()), 0)
+        self.assertTrue(g.get())
+        g = QgsGeometry.fromWkt('MultiLineString((0 0, 1 1),(13 2, 14 1))')
+        self.assertEqual(len(g.get()), 2)
+        self.assertTrue(g.get())
+        self.assertEqual(g.get().geometryN(0).asWkt(), 'LineString (0 0, 1 1)')
+        self.assertEqual(g.get().geometryN(1).asWkt(), 'LineString (13 2, 14 1)')
+        with self.assertRaises(IndexError):
+            g.get().geometryN(-1)
+        with self.assertRaises(IndexError):
+            g.get().geometryN(2)
+
     def testIntersection(self):
         myLine = QgsGeometry.fromPolylineXY([
             QgsPointXY(0, 0),
