diff --git a/opencage/batch.py b/opencage/batch.py index 7af9e0f..1e3b74e 100644 --- a/opencage/batch.py +++ b/opencage/batch.py @@ -213,11 +213,13 @@ async def write_one_geocoding_result(self, csv_output, row_id, geocoding_result, if geocoding_result is None: row.append('') elif column in geocoding_result: - row.append(geocoding_result[column]) + row.append(self.deep_get_result_value(geocoding_result, [column], '')) elif column in geocoding_result['components']: - row.append(geocoding_result['components'][column]) + row.append(self.deep_get_result_value(geocoding_result, ['components', column], '')) elif column in geocoding_result['geometry']: - row.append(geocoding_result['geometry'][column]) + row.append(self.deep_get_result_value(geocoding_result, ['geometry', column], '')) + elif column == 'FIPS': + row.append(self.deep_get_result_value(geocoding_result, ['annotations', 'FIPS', 'county'], '')) else: row.append('') @@ -235,7 +237,14 @@ async def write_one_geocoding_result(self, csv_output, row_id, geocoding_result, csv_output.writerow(row) self.write_counter = self.write_counter + 1 - def log(self, message): if not self.options.quiet: sys.stderr.write(f"{message}\n") + + def deep_get_result_value(self, data, keys, default=None): + for key in keys: + if isinstance(data, dict): + data = data.get(key, default) + else: + return default + return data diff --git a/test/test_batch.py b/test/test_batch.py new file mode 100644 index 0000000..e9484ba --- /dev/null +++ b/test/test_batch.py @@ -0,0 +1,24 @@ +from opencage.batch import OpenCageBatchGeocoder + +batch = OpenCageBatchGeocoder({}) + +def test_deep_get_result_value(): + result = { + 'annotations': { + 'FIPS': { + 'state': 'CA' + } + }, + 'components': { + 'street': 'Main Road' + } + } + + assert batch.deep_get_result_value(result, ['hello', 'world']) == None + + assert batch.deep_get_result_value(result, ['components', 'street']) == 'Main Road' + assert batch.deep_get_result_value(result, ['components', 'city']) == None + assert batch.deep_get_result_value(result, ['components', 'city'], '') == '' + + assert batch.deep_get_result_value([], ['hello', 'world']) == None + assert batch.deep_get_result_value(None, ['hello', 'world']) == None