Merge pull request #49 from johntruckenbrodt/feature/vector_addfield_…

…values

Vector.addfield: new argument 'values'

spatialist/tests/test_spatial.py

@@ -4,13 +4,13 @@
 import platform
 import numpy as np
 from osgeo import ogr, gdal
-from spatialist import crsConvert, haversine, Raster, stack, ogr2ogr, gdal_translate, gdal_rasterize, bbox, rasterize, \
-    gdalwarp, utm_autodetect, coordinate_reproject, cmap_mpl2gdal
-from spatialist.raster import Dtype, png
-from spatialist.vector import feature2vector, dissolve, Vector, intersect
+from spatialist.raster import Dtype, png, Raster, stack, rasterize
+from spatialist.vector import feature2vector, dissolve, Vector, intersect, bbox
 from spatialist.envi import hdr, HDRobject
 from spatialist.sqlite_util import sqlite_setup, __Handler
 from spatialist.ancillary import parallel_apply_along_axis
+from spatialist.auxil import (crsConvert, haversine, ogr2ogr, gdal_translate, gdal_rasterize, gdalwarp,
+                              utm_autodetect, coordinate_reproject, cmap_mpl2gdal)
 
 import logging
 
@@ -426,3 +426,17 @@ def test_png(tmpdir, testdata):
     outname = os.path.join(str(tmpdir), 'test_rgb.png')
     with Raster(src) as ras:
         png(src=ras, dst=outname, percent=100, scale=(2, 98), worldfile=True)
+
+
+def test_addfield():
+    extent = {'xmin': 10, 'xmax': 11, 'ymin': 50, 'ymax': 51}
+    with bbox(coordinates=extent, crs=4326) as box:
+        box.addfield(name='test1', type=ogr.OFTString, values=['a'])
+        box.addfield(name='test2', type=ogr.OFTStringList, values=[['a', 'b']])
+        box.addfield(name='test3', type=ogr.OFTInteger, values=[1])
+        box.addfield(name='test4', type=ogr.OFTIntegerList, values=[[1, 2]])
+        box.addfield(name='test5', type=ogr.OFTInteger64, values=[1])
+        box.addfield(name='test6', type=ogr.OFTInteger64List, values=[[1, 2]])
+        box.addfield(name='test7', type=ogr.OFTReal, values=[1])
+        box.addfield(name='test8', type=ogr.OFTRealList, values=[[1., 2.]])
+        box.addfield(name='test9', type=ogr.OFTBinary, values=[b'1'])

spatialist/vector.py

@@ -168,7 +168,7 @@ def addfeature(self, geometry, fields=None):
         feature = None
         self.init_features()
 
-    def addfield(self, name, type, width=10):
+    def addfield(self, name, type, width=10, values=None):
         """
         add a field to the vector layer
 
@@ -178,18 +178,39 @@ def addfield(self, name, type, width=10):
             the field name
         type: int
             the OGR Field Type (OFT), e.g. ogr.OFTString.
-            See `Module ogr <https://gdal.org/python/osgeo.ogr-module.html>`_.
+            See :class:`osgeo.ogr.FieldDefn`.
         width: int
             the width of the new field (only for ogr.OFTString fields)
+        values: list
+            an optional list with values for each feature to assign to the new field.
+            The length must be identical to the number of features.
 
         Returns
         -------
 
         """
-        fieldDefn = ogr.FieldDefn(name, type)
+        type_name = ogr.GetFieldTypeName(type)
+        field_defn = ogr.FieldDefn(name, type)
         if type == ogr.OFTString:
-            fieldDefn.SetWidth(width)
-        self.layer.CreateField(fieldDefn)
+            field_defn.SetWidth(width)
+        self.layer.CreateField(field_defn)
+        if type_name in ['String', 'Integer', 'Real', 'Binary']:
+            method_name = 'SetField'
+        elif type_name in ['StringList', 'DoubleList', 'IntegerList',
+                           'Integer64', 'Integer64List']:
+            method_name = f'SetField{type_name}'
+        elif type_name == 'RealList':
+            method_name = 'SetFieldDoubleList'
+        else:
+            raise ValueError(f'Unsupported field type: {type_name}')
+        if values is not None:
+            if len(values) != self.nfeatures:
+                raise RuntimeError('number of values does not match number of features')
+            for i, feature in enumerate(self.layer):
+                index = feature.GetFieldIndex(name)
+                method = getattr(feature, method_name)
+                method(index, values[i])
+                self.layer.SetFeature(feature)
 
     def addlayer(self, name, srs, geomType):
         """