[FEATURE] Add with edit(layer): to python

m-kuhn · m-kuhn · commit 35fc2902f332 · 2015-08-12T13:41:23.000+02:00
Example:

    with edit(layer):
        f=layer.getFeatures().next()
        f[0]=5
        layer.updateFeature(f)

This will automatically call commitChanges() in the end.
If any exception occurs, it will rollBack() all the changes.
diff --git a/python/__init__.py b/python/__init__.py
@@ -72,6 +72,27 @@ def __hash__(self):
 except ImportError:
     pass
 
+# Define a `with edit(layer)` statement
+
+class edit:
+    def __init__(self,layer):
+        self.layer = layer
+
+    def __enter__(self):
+        assert self.layer.startEditing()
+        return self.layer
+
+    def __exit__(self, ex_type, ex_value, traceback):
+        if ex_type is None:
+            assert self.layer.commitChanges()
+            return True
+        else:
+            self.layer.rollBack()
+            return False
+
+from qgis import core
+core.edit = edit
+
 
 def mapping_feature(feature):
     geom = feature.geometry()
diff --git a/tests/src/python/CMakeLists.txt b/tests/src/python/CMakeLists.txt
@@ -52,6 +52,7 @@ ADD_PYTHON_TEST(PyQgsSpatialiteProvider test_provider_spatialite.py)
 ADD_PYTHON_TEST(PyQgsShapefileProvider test_provider_shapefile.py)
 ADD_PYTHON_TEST(PyQgsMemoryProvider test_provider_memory.py)
 ADD_PYTHON_TEST(PyQgsVectorColorRamp test_qgsvectorcolorramp.py)
+ADD_PYTHON_TEST(PyQgsSyntacticSugar test_syntactic_sugar.py)
 
 IF (NOT WIN32)
   ADD_PYTHON_TEST(PyQgsLogger test_qgslogger.py)
diff --git a/tests/src/python/test_syntactic_sugar.py b/tests/src/python/test_syntactic_sugar.py
@@ -0,0 +1,76 @@
+# -*- coding: utf-8 -*-
+"""QGIS Unit tests for some syntactic sugar in python
+
+.. note:: This program is free software; you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation; either version 2 of the License, or
+(at your option) any later version.
+"""
+__author__ = 'Matthias Kuhn'
+__date__ = '12.8.2015'
+__copyright__ = 'Copyright 2015, The QGIS Project'
+# This will get replaced with a git SHA1 when you do a git archive
+__revision__ = '$Format:%H$'
+
+import qgis
+
+from utilities import (unittest,
+                       TestCase,
+                       getQgisTestApp
+                      )
+from qgis.core import (edit,
+                       QgsFeature,
+                       QgsGeometry,
+                       QgsVectorLayer
+                      )
+
+getQgisTestApp()
+
+class TestSyntacticSugar(TestCase):
+
+    def testEdit(self):
+        """Test `with edit(layer):` code"""
+
+        ml=QgsVectorLayer("Point?crs=epsg:4236&field=id:integer&field=value:double",
+                          "test_data", "memory")
+        # Data as list of x, y, id, value
+        assert ml.isValid()
+        fields=ml.fields()
+
+        # Check insert
+        with edit(ml):
+            feat=QgsFeature(fields)
+            feat['id']=1
+            feat['value']=0.9
+            assert ml.addFeature(feat)
+
+        assert ml.dataProvider().getFeatures().next()['value']==0.9
+
+        # Check update
+        with edit(ml):
+            f = ml.getFeatures().next()
+            f['value']=9.9
+            assert ml.updateFeature(f)
+
+        assert ml.dataProvider().getFeatures().next()['value']==9.9
+
+        # Check for rollBack after exceptions
+        with self.assertRaises(NameError):
+            with edit(ml):
+                f = ml.getFeatures().next()
+                f['value']=3.8
+                crashycrash()
+
+        assert ml.dataProvider().getFeatures().next()['value']==9.9
+        assert ml.getFeatures().next()['value']==9.9
+
+        # Check for `as`
+        with edit(ml) as l:
+            f = l.getFeatures().next()
+            f['value']=10
+            assert l.updateFeature(f)
+
+        assert ml.dataProvider().getFeatures().next()['value']==10
+
+if __name__ == "__main__":
+        unittest.main()
