Merge pull request #15 from eyeseast/label-results

Label results with geocoder used

Author: eyeseast

Diff for: geocode_sqlite/cli.py

@@ -85,6 +85,10 @@ def geocode(ctx, geocoder):
         click.echo(f"Adding column: {longitude}")
         table.add_column(longitude, float)
 
+    if "geocoder" not in table.columns_dict:
+        click.echo("Adding geocoder column")
+        table.add_column("geocoder", str)
+
     # always use a rate limiter, even if delay is zero
     geocode = RateLimiter(geocoder.geocode, min_delay_seconds=delay)
 
@@ -93,13 +97,12 @@ def geocode(ctx, geocoder):
     )
 
     done = 0
+    errors = []
 
     gen = geocode_list(
         rows, geocode, location, latitude_column=latitude, longitude_column=longitude
     )
 
-    errors = []
-
     with click.progressbar(gen, length=count, label=f"{count} rows") as bar:
         for row, success in bar:
             pks = [row[pk] for pk in table.pks]

Diff for: geocode_sqlite/utils.py

@@ -2,7 +2,6 @@
 This is the Python interface
 """
 import logging
-from geopy import geocoders
 from geopy.extra.rate_limiter import RateLimiter
 from sqlite_utils import Database
 
@@ -39,31 +38,41 @@ def geocode_table(
     table = db[table_name]
 
     if latitude_column not in table.columns_dict:
+        log.info(f"Adding latitude column: {latitude_column}")
         table.add_column(latitude_column, float)
 
     if longitude_column not in table.columns_dict:
+        log.info(f"Adding longitude column: {longitude_column}")
         table.add_column(longitude_column, float)
 
-    if force:
-        rows = table.rows
-    else:
-        rows = table.rows_where(
-            f"{latitude_column} IS NULL OR {longitude_column} IS NULL"
-        )
+    if "geocoder" not in table.columns_dict:
+        log.info("Adding geocoder column")
+        table.add_column("geocoder", str)
+
+    rows, todo = select_ungeocoded(
+        db,
+        table,
+        latitude_column=latitude_column,
+        longitude_column=longitude_column,
+        force=force,
+    )
 
-    if delay:
-        geocode = RateLimiter(geocoder.geocode, min_delay_seconds=delay)
-    else:
-        geocode = geocoder.geocode
+    # always use a rate limiter, even with no delay
+    geocode = RateLimiter(geocoder.geocode, min_delay_seconds=delay)
 
     count = 0
+    log.info(f"Geocoding {todo} rows from {table.name}")
     for row in rows:
         result = geocode_row(geocode, query_template, row)
         if result:
             pks = [row[pk] for pk in table.pks]
             table.update(
                 pks,
-                {latitude_column: result.latitude, longitude_column: result.longitude},
+                {
+                    latitude_column: result.latitude,
+                    longitude_column: result.longitude,
+                    "geocoder": geocoder.__class__.__name__,
+                },
             )
             count += 1
 
@@ -95,6 +104,7 @@ def geocode_list(
         if result:
             row[longitude_column] = result.longitude
             row[latitude_column] = result.latitude
+            row["geocoder"] = get_geocoder_class(geocode)
 
         yield row, bool(result)
 
@@ -126,3 +136,13 @@ def select_ungeocoded(
     rows = table.rows_where(f"{latitude_column} IS NULL OR {longitude_column} IS NULL")
 
     return rows, count
+
+
+def get_geocoder_class(geocode):
+    "Walk back up to the original geocoder class"
+
+    if isinstance(geocode, RateLimiter):
+        return geocode.func.__self__.__class__.__name__
+
+    # unwrapped function
+    return geocode.__self__.__class__.__name__

Diff for: tests/test_geocode_sqlite.py

@@ -231,3 +231,14 @@ def test_geocode_list(db, geocoder):
     for row, success in done:
         assert success
         assert row == table.get(row["id"])
+
+
+def test_label_results(db, geocoder):
+    table = db[TABLE_NAME]
+
+    # geocode it once
+    geocode_table(db, TABLE_NAME, geocoder, "{id}")
+
+    for row in table.rows:
+        assert "geocoder" in row
+        assert row["geocoder"] == geocoder.__class__.__name__