Coverage for gws-app/gws/base/exporter/util.py: 90%

1"""Helpers for exporters.""" 

2 

3from typing import Optional 

4import gws 

5import gws.lib.gdalx 

6import gws.lib.mime 

7import gws.lib.osx 

8import gws.lib.zipx 

9 

10 

11class Group(gws.Data): 

12 """Export-ready features of one model.""" 

13 

14 title: str 

15 """Group title: the model title, table name or model uid.""" 

16 records: list[gws.FeatureRecord] 

17 """Records with the exported attributes and the shape.""" 

18 columns: dict[str, gws.AttributeType] 

19 """Exported columns and their types.""" 

20 geomType: Optional[gws.GeometryType] 

21 """Geometry type, taken from the first feature with a shape.""" 

22 crs: Optional[gws.Crs] 

23 """CRS, taken from the first feature with a shape.""" 

24 

25 

26def group_features(ea: gws.ExportArgs, er: gws.ExportResult) -> list[Group]: 

27 """Group features by model and determine the export columns and geometry type. 

28 

29 Columns are the model fields with an attribute type supported by the 

30 exporter (or by GDAL, if the exporter does not restrict the types). 

31 Features that do not fit their group are skipped and reported in 

32 ``er.errors``, see ``gws.base.exporter``. 

33 

34 Args: 

35 ea: Export arguments. 

36 er: Export result, used to report errors. 

37 

38 Returns: 

39 A list of groups, one per model that has exportable features. 

40 """ 

41 

42 if not ea.features: 

43 return [] 

44 

45 by_model: dict[str, list[gws.Feature]] = {} 

46 for f in ea.features: 

47 by_model.setdefault(f.model.uid, []).append(f) 

48 

49 ls = [] 

50 

51 for _, features in by_model.items(): 

52 grp = _create_group(features, ea, er) 

53 if grp: 

54 ls.append(grp) 

55 

56 return ls 

57 

58 

59def _create_group(features: list[gws.Feature], ea: gws.ExportArgs, er: gws.ExportResult) -> Optional[Group]: 

60 """Create a group for the features of one model, or ``None`` if no feature fits.""" 

61 grp = Group( 

62 records=[], 

63 title='', 

64 columns={}, 

65 geomType=None, 

66 srid=0, 

67 ) 

68 

69 f = features[0] 

70 

71 grp.title = f.model.title 

72 if not grp.title and hasattr(f.model, 'tableName'): 

73 grp.title = getattr(f.model, 'tableName').split('.')[-1] 

74 if not grp.title: 

75 grp.title = f.model.uid 

76 

77 types = ea.exporter.supportedAttributeTypes or gws.lib.gdalx.supported_attribute_types() 

78 for fld in f.model.fields: 

79 if fld.attributeType in types: 

80 grp.columns[fld.name] = fld.attributeType 

81 

82 for n, f in enumerate(features, 1): 

83 if ea.notify: 

84 ea.notify('') 

85 rec = _feature_to_record(f, grp, ea, er) 

86 if rec: 

87 grp.records.append(rec) 

88 

89 return grp if grp.records else None 

90 

91 

92def _feature_to_record(f: gws.Feature, grp: Group, ea: gws.ExportArgs, er: gws.ExportResult) -> Optional[gws.FeatureRecord]: 

93 """Convert a feature to a record of the group, or report an error and return ``None``.""" 

94 uid = f.uid() 

95 

96 sh = f.shape() 

97 

98 if sh: 

99 if grp.geomType and grp.geomType != sh.type and not ea.exporter.withMixedGeometry: 

100 if len(er.errors) < ea.maxErrors: 

101 er.errors.append(f'{uid}: inconsistent geometry: {grp.geomType} and {sh.type}') 

102 return 

103 

104 if grp.crs and grp.crs != sh.crs and not ea.exporter.withMixedCrs: 

105 if len(er.errors) < ea.maxErrors: 

106 er.errors.append(f'{uid}: inconsistent CRS: {grp.crs} and {sh.crs}') 

107 return 

108 

109 if not grp.geomType: 

110 grp.geomType = sh.type 

111 grp.crs = sh.crs 

112 

113 elif not ea.exporter.withNoGeometry: 

114 if len(er.errors) < ea.maxErrors: 

115 er.errors.append(f'{uid}: has no geometry') 

116 return 

117 

118 return gws.FeatureRecord( 

119 attributes={c: f.get(c) for c in grp.columns}, 

120 shape=sh, 

121 ) 

122 

123 

124## 

125 

126 

127def run_gdal_vector_export(driver_name: str, mime_type: str, ea: gws.ExportArgs, er: gws.ExportResult): 

128 """Export features with a GDAL vector driver. 

129 

130 Writes one file per group, or a single file with one layer per group if 

131 the exporter has ``withMultiLayer``. Several files are zipped. Sets the 

132 path, mime type and counts in the export result. Does nothing if there 

133 are no features to export. 

134 

135 Args: 

136 driver_name: GDAL driver name, e.g. ``GeoJSON``. 

137 mime_type: Mime type of the result file. 

138 ea: Export arguments. 

139 er: Export result, filled by this function. 

140 

141 Raises: 

142 ``gws.Error``: If the driver is not supported. 

143 """ 

144 

145 di = gws.lib.gdalx.get_driver(driver_name) 

146 if not di: 

147 raise gws.Error(f'unsupported driver: {driver_name}') 

148 

149 groups = group_features(ea, er) 

150 if not groups: 

151 return 

152 

153 base_dir = gws.u.ephemeral_dir(gws.u.random_string(64)) 

154 ext = di.extensions[0] if di.extensions else 'dat' 

155 

156 if ea.exporter.withMultiLayer: 

157 path = base_dir + f'/export.{ext}' 

158 with gws.lib.gdalx.open_vector(path, 'w', driver=di.name, options=ea.exporter.options) as ds: 

159 for grp in groups: 

160 la = ds.create_layer(grp.title, grp.columns, grp.geomType, grp.crs) 

161 fids = la.insert(grp.records) 

162 er.numFeaturesExported += len(fids) 

163 else: 

164 for grp in groups: 

165 path = base_dir + f'/{gws.u.to_uid(grp.title)}.{ext}' 

166 with gws.lib.gdalx.open_vector(path, 'w', driver=di.name, options=ea.exporter.options) as ds: 

167 la = ds.create_layer(grp.title, grp.columns, grp.geomType, grp.crs) 

168 fids = la.insert(grp.records) 

169 er.numFeaturesExported += len(fids) 

170 

171 paths = list(gws.lib.osx.find_files(base_dir, deep=False)) 

172 er.numFiles = len(paths) 

173 if er.numFiles > 1: 

174 er.path = base_dir + '/export.zip' 

175 er.mimeType = gws.lib.mime.ZIP 

176 gws.lib.zipx.zip_to_path(er.path, paths, flat=True) 

177 else: 

178 er.path = paths[0] 

179 er.mimeType = mime_type