Coverage for py_tools_ds/io/vector/writer.py: 0%

1# -*- coding: utf-8 -*- 

2 

3# py_tools_ds - A collection of geospatial data analysis tools that simplify standard 

4# operations when handling geospatial raster and vector data as well as projections. 

5# 

6# Copyright (C) 2016-2021 

7# - Daniel Scheffler (GFZ Potsdam, daniel.scheffler@gfz-potsdam.de) 

8# - Helmholtz Centre Potsdam - GFZ German Research Centre for Geosciences Potsdam, 

9# Germany (https://www.gfz-potsdam.de/) 

10# 

11# This software was developed within the context of the GeoMultiSens project funded 

12# by the German Federal Ministry of Education and Research 

13# (project grant code: 01 IS 14 010 A-C). 

14# 

15# Licensed under the Apache License, Version 2.0 (the "License"); 

16# you may not use this file except in compliance with the License. 

17# You may obtain a copy of the License at 

18# 

19# http://www.apache.org/licenses/LICENSE-2.0 

20# 

21# Unless required by applicable law or agreed to in writing, software 

22# distributed under the License is distributed on an "AS IS" BASIS, 

23# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

24# See the License for the specific language governing permissions and 

25# limitations under the License. 

26 

27import os 

28 

29from osgeo import ogr, osr 

30 

31from ...dtypes.conversion import get_dtypeStr 

32from ...geo.projection import EPSG2WKT 

33 

34 

35__author__ = "Daniel Scheffler" 

36 

37 

38def write_shp(path_out, shapely_geom, prj=None, attrDict=None): 

39 shapely_geom = [shapely_geom] if not isinstance(shapely_geom, list) else shapely_geom 

40 attrDict = [attrDict] if not isinstance(attrDict, list) else attrDict 

41 # print(len(shapely_geom)) 

42 # print(len(attrDict)) 

43 assert len(shapely_geom) == len(attrDict), "'shapely_geom' and 'attrDict' must have the same length." 

44 assert os.path.exists(os.path.dirname(path_out)), 'Directory %s does not exist.' % os.path.dirname(path_out) 

45 

46 print('Writing %s ...' % path_out) 

47 if os.path.exists(path_out): 

48 os.remove(path_out) 

49 ds = ogr.GetDriverByName("Esri Shapefile").CreateDataSource(path_out) 

50 

51 if prj is not None: 

52 prj = prj if not isinstance(prj, int) else EPSG2WKT(prj) 

53 srs = osr.SpatialReference() 

54 srs.ImportFromWkt(prj) 

55 else: 

56 srs = None 

57 

58 geom_type = list(set([gm.type for gm in shapely_geom])) 

59 assert len(geom_type) == 1, 'All shapely geometries must belong to the same type. Got %s.' % geom_type 

60 

61 layer = \ 

62 ds.CreateLayer('', srs, ogr.wkbPoint) if geom_type[0] == 'Point' else\ 

63 ds.CreateLayer('', srs, ogr.wkbLineString) if geom_type[0] == 'LineString' else \ 

64 ds.CreateLayer('', srs, ogr.wkbPolygon) if geom_type[0] == 'Polygon' else None # FIXME 

65 

66 if isinstance(attrDict[0], dict): 

67 for attr in attrDict[0].keys(): 

68 assert len(attr) <= 10, "ogr does not support fieldnames longer than 10 digits. '%s' is too long" % attr 

69 DTypeStr = get_dtypeStr(attrDict[0][attr]) 

70 FieldType = \ 

71 ogr.OFTInteger if DTypeStr.startswith('int') else \ 

72 ogr.OFTReal if DTypeStr.startswith('float') else \ 

73 ogr.OFTString if DTypeStr.startswith('str') else \ 

74 ogr.OFTDateTime if DTypeStr.startswith('date') else None 

75 FieldDefn = ogr.FieldDefn(attr, FieldType) 

76 if DTypeStr.startswith('float'): 

77 FieldDefn.SetPrecision(6) 

78 layer.CreateField(FieldDefn) # Add one attribute 

79 

80 for i in range(len(shapely_geom)): 

81 # Create a new feature (attribute and geometry) 

82 feat = ogr.Feature(layer.GetLayerDefn()) 

83 feat.SetGeometry(ogr.CreateGeometryFromWkb(shapely_geom[i].wkb)) # Make a geometry, from Shapely object 

84 

85 list_attr2set = attrDict[0].keys() if isinstance(attrDict[0], dict) else [] 

86 

87 for attr in list_attr2set: 

88 val = attrDict[i][attr] 

89 DTypeStr = get_dtypeStr(val) 

90 val = int(val) if DTypeStr.startswith('int') else float(val) if DTypeStr.startswith('float') else \ 

91 str(val) if DTypeStr.startswith('str') else val 

92 feat.SetField(attr, val) 

93 

94 layer.CreateFeature(feat) 

95 feat.Destroy() 

96 

97 # Save and close everything 

98 del ds, layer