more robust fetch values from result data structure

OpenCageData · Sep 25, 2024 · 4ee6787 · 4ee6787
1 parent a49a1ef
commit 4ee6787
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 4 deletions.
diff --git 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
@@ -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