Skip to content

Commit 21f578d

Browse files
iferencikiferencik
authored
fix admin cgaz postprocessing (#515)
1 parent 8e6a1d4 commit 21f578d

1 file changed

Lines changed: 72 additions & 38 deletions

File tree

‎rapida/admin/cgaz.py‎

Lines changed: 72 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -61,57 +61,91 @@ def fetch_admin(bbox=None, admin_level=None, clip=False, destination_path=None,
6161
progress.remove_task(translate_task)
6262

6363
progress.update(task, advance=1, description="[green]Downloaded admin data", refresh=True)
64-
with gdal.OpenEx(destination_path, gdal.OF_VECTOR|gdal.OF_UPDATE) as ds:
64+
with gdal.OpenEx(destination_path, gdal.OF_VECTOR | gdal.OF_UPDATE) as ds:
6565
layer = ds.GetLayerByName(dst_layer_name or f'admin{admin_level}')
66+
67+
# Guard against empty results
68+
if layer is None:
69+
logger.error(f"Layer not found. BBOX {bbox} likely resulted in 0 features.")
70+
progress.update(task, description="[red]No features found in BBOX", refresh=True)
71+
ds = None
72+
return None
73+
74+
# Define feature_count here
75+
feature_count = layer.GetFeatureCount()
76+
6677
iso3_field_index = layer.GetLayerDefn().GetFieldIndex("iso3")
6778

68-
if iso3_field_index >=0:
79+
if iso3_field_index >= 0:
6980
countries = set([f.GetField("iso3") for f in layer])
7081
if len(countries) == 0:
7182
raise Exception(f"No countries were found in {bbox}")
83+
layer.ResetReading()
7284
else:
7385
raise Exception(f'{url} does not contain iso3 field!')
7486

87+
# --- 1. BATCH ALL SCHEMA ALTERATIONS FIRST ---
7588
h3id_field_index = layer.GetLayerDefn().GetFieldIndex("h3id")
76-
if h3id_field_index < 0:
77-
h3id_field = ogr.FieldDefn("h3id", ogr.OFTInteger64)
78-
layer.CreateField(h3id_field)
79-
for feature in layer:
80-
geom = feature.GetGeometryRef()
81-
centroid = geom.Centroid()
82-
h3id = h3.latlng_to_cell(lat=centroid.GetY(), lng=centroid.GetX(),
83-
res=h3id_precision)
84-
feature.SetField("h3id", h3id)
85-
layer.SetFeature(feature)
86-
name_field_index = layer.GetLayerDefn().GetFieldIndex("name")
87-
source_name_field = f"admin{admin_level}_name"
88-
source_name_index = layer.GetLayerDefn().GetFieldIndex(source_name_field)
89+
needs_h3 = h3id_field_index < 0
90+
if needs_h3:
91+
layer.CreateField(ogr.FieldDefn("h3id", ogr.OFTInteger64))
8992

90-
if source_name_index < 0:
93+
source_name_field = f"admin{admin_level}_name"
94+
if layer.GetLayerDefn().GetFieldIndex(source_name_field) < 0:
9195
raise Exception(f"Source {url} does not contain expected name field: {source_name_field}")
9296

93-
if name_field_index < 0:
94-
name_field = ogr.FieldDefn("name", ogr.OFTString)
95-
layer.CreateField(name_field)
96-
name_field_index = layer.GetLayerDefn().GetFieldIndex("name")
97-
if name_field_index < 0:
98-
raise Exception("Error: 'name' field was not created!")
99-
layer.ResetReading()
100-
for feature in layer:
101-
fid = feature.GetFID()
102-
name_value = feature.GetField(source_name_field)
103-
feature.SetField("name", name_value)
104-
if layer.SetFeature(feature) != 0:
105-
raise Exception(f"Failed to update feature {fid} with name '{name_value}'")
106-
107-
if not keep_disputed_areas:
108-
layer.ResetReading()
109-
for feature in layer:
110-
iso3_country_code = feature.GetField("iso3")
111-
if iso3_country_code not in COUNTRY_CODES:
112-
layer.DeleteFeature(feature.GetFID())
97+
name_field_index = layer.GetLayerDefn().GetFieldIndex("name")
98+
needs_name = name_field_index < 0
99+
if needs_name:
100+
layer.CreateField(ogr.FieldDefn("name", ogr.OFTString))
101+
102+
# --- 2. SINGLE PASS FOR ALL ROW UPDATES ---
103+
layer.ResetReading()
104+
fids_to_delete = []
105+
106+
update_task = progress.add_task("[yellow]Updating features...", total=feature_count)
107+
108+
for feature in layer:
109+
update_needed = False
110+
111+
if not keep_disputed_areas and feature.GetField("iso3") not in COUNTRY_CODES:
112+
fids_to_delete.append(feature.GetFID())
113+
else:
114+
if needs_h3:
115+
geom = feature.GetGeometryRef()
116+
if geom is not None and not geom.IsEmpty():
117+
centroid = geom.Centroid()
118+
if centroid is not None:
119+
h3id = h3.latlng_to_cell(lat=centroid.GetY(), lng=centroid.GetX(),
120+
res=h3id_precision)
121+
feature.SetField("h3id", h3id)
122+
update_needed = True
123+
124+
if needs_name:
125+
name_value = feature.GetField(source_name_field)
126+
safe_name = str(name_value) if name_value is not None else ""
127+
feature.SetField("name", safe_name)
128+
update_needed = True
129+
130+
if update_needed:
131+
if layer.SetFeature(feature) != 0:
132+
logger.warning(f"Failed to update feature {feature.GetFID()}")
133+
134+
progress.advance(update_task)
135+
136+
progress.remove_task(update_task)
137+
138+
# --- 3. EXECUTE DELETIONS ---
139+
if fids_to_delete:
140+
del_task = progress.add_task("[red]Deleting disputed areas...", total=len(fids_to_delete))
141+
for fid in fids_to_delete:
142+
layer.DeleteFeature(fid)
143+
progress.advance(del_task)
144+
progress.remove_task(del_task)
113145

114146
ds.FlushCache()
115-
progress.update(task, advance=1, description="[green]Download Completed", refresh=True)
116-
return None
147+
148+
149+
progress.update(task, advance=1, description="[green]Download completed.", refresh=True)
150+
return None
117151

0 commit comments

Comments
 (0)