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

1"""Export utilities.""" 

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 title: str 

13 records: list[gws.FeatureRecord] 

14 columns: dict[str, gws.AttributeType] 

15 geomType: Optional[gws.GeometryType] 

16 crs: Optional[gws.Crs] 

17 

18 

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

20 """Group features by model, determine export columns and geometry type. 

21 

22 Args: 

23 ea: Export arguments. 

24 er: Export result, used to report errors and counts. 

25 

26 Returns: 

27 List of groups, one per model, with export-ready records. 

28 """ 

29 

30 if not ea.features: 

31 return [] 

32 

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

34 for f in ea.features: 

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

36 

37 ls = [] 

38 

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

40 grp = _create_group(features, ea, er) 

41 if grp: 

42 ls.append(grp) 

43 

44 return ls 

45 

46 

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

48 grp = Group( 

49 records=[], 

50 title='', 

51 columns={}, 

52 geomType=None, 

53 srid=0, 

54 ) 

55 

56 f = features[0] 

57 

58 grp.title = f.model.title 

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

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

61 if not grp.title: 

62 grp.title = f.model.uid 

63 

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

65 for fld in f.model.fields: 

66 if fld.attributeType in types: 

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

68 

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

70 if ea.notify: 

71 ea.notify('') 

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

73 if rec: 

74 grp.records.append(rec) 

75 

76 return grp if grp.records else None 

77 

78 

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

80 uid = f.uid() 

81 

82 sh = f.shape() 

83 

84 if sh: 

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

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

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

88 return 

89 

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

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

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

93 return 

94 

95 if not grp.geomType: 

96 grp.geomType = sh.type 

97 grp.crs = sh.crs 

98 

99 elif not ea.exporter.withNoGeometry: 

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

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

102 return 

103 

104 return gws.FeatureRecord( 

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

106 shape=sh, 

107 ) 

108 

109 

110## 

111 

112 

113def run_gdal_vector_export(driver_name: str, mime: str, ea: gws.ExportArgs, er: gws.ExportResult): 

114 """Run the export for a GDAL vector driver.""" 

115 

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

117 if not di: 

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

119 

120 groups = group_features(ea, er) 

121 if not groups: 

122 return 

123 

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

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

126 

127 if ea.exporter.withMultiLayer: 

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

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

130 for grp in groups: 

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

132 fids = la.insert(grp.records) 

133 er.numFeaturesExported += len(fids) 

134 else: 

135 for grp in groups: 

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

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

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

139 fids = la.insert(grp.records) 

140 er.numFeaturesExported += len(fids) 

141 

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

143 er.numFiles = len(paths) 

144 if er.numFiles > 1: 

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

146 er.mime = gws.lib.mime.ZIP 

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

148 else: 

149 er.path = paths[0] 

150 er.mime = mime