diff --git a/openaddr/cache.py b/openaddr/cache.py index c3c74916..bc18c360 100644 --- a/openaddr/cache.py +++ b/openaddr/cache.py @@ -28,6 +28,7 @@ from .conform import GEOM_FIELDNAME from . import util +from . import progress def mkdirsp(path): try: @@ -298,10 +299,16 @@ def download(self, source_urls, workdir, source_config): raise DownloadError('{} response from {}'.format(resp.status_code, source_url)) size = 0 + content_length = resp.headers.get('Content-Length') + progress.reporter.start_phase( + 'Downloading', total=int(content_length) if content_length and content_length.isdigit() else None, + unit='bytes') with open(file_path, 'wb') as fp: for chunk in resp.iter_content(self.CHUNK): size += len(chunk) fp.write(chunk) + progress.reporter.advance(len(chunk)) + progress.reporter.end_phase() output_files.append(file_path) @@ -414,17 +421,21 @@ def download(self, source_urls, workdir, source_config): field_names.append(GEOM_FIELDNAME) # Get the count of rows in the layer + row_count = None try: row_count = downloader.get_feature_count() _L.info("Source has {} rows".format(row_count)) except EsriDownloadError: _L.info("Source doesn't support count") + progress.reporter.start_phase('Downloading from ArcGIS', total=row_count) + with open(file_path, 'w', encoding='utf-8') as f: writer = csv.DictWriter(f, fieldnames=field_names) writer.writeheader() for feature in downloader: + point = None try: geom = feature.get('geometry') or {} row = feature.get('properties') or {} @@ -453,12 +464,21 @@ def download(self, source_urls, workdir, source_config): raise TypeError("Geometry has non-finite coordinates") shp = shape(geom) + point = None if shp.is_empty else (shp if shp.geom_type == 'Point' else shp.representative_point()) row[GEOM_FIELDNAME] = shp.wkt writer.writerow({fn: row.get(fn) for fn in field_names}) size += 1 + progress.reporter.advance() + if point is not None: + progress.reporter.cell(point.x, point.y, 'ok') except TypeError: _L.debug("Skipping a geometry", exc_info=True) + progress.reporter.count('skipped') + if point is not None: + progress.reporter.cell(point.x, point.y, 'error') + + progress.reporter.end_phase() _L.info("Downloaded %s ESRI features for file %s", size, file_path) output_files.append(file_path) diff --git a/openaddr/conform.py b/openaddr/conform.py index efd3d6db..11d3e07a 100644 --- a/openaddr/conform.py +++ b/openaddr/conform.py @@ -15,10 +15,11 @@ import osgeo import shutil +from . import progress from .geojson import stream_geojson from shapely.wkt import loads as wkt_loads -from shapely.geometry import mapping +from shapely.geometry import mapping, shape from zipfile import ZipFile from locale import getpreferredencoding @@ -230,6 +231,7 @@ def decompress(self, source_paths, workdir, filenames): # Extract contents of zip file into expand_path directory. found = set() + progress.reporter.start_phase('Unzipping', total=self._extract_size(source_paths, filenames), unit='bytes') for source_path in source_paths: found |= self._extract_zip(source_path, expand_path, filenames) @@ -261,6 +263,7 @@ def fully_satisfied(): if zip_path in processed: continue processed.add(zip_path) + progress.reporter.add_total(self._extract_size([zip_path], filenames)) found |= self._extract_zip(zip_path, os.path.dirname(zip_path), filenames) if fully_satisfied(): break @@ -281,8 +284,20 @@ def fully_satisfied(): output_files.append(os.path.join(dirpath, filename)) _L.debug("Expanded file {}".format(output_files[-1])) + progress.reporter.end_phase() return output_files + @staticmethod + def _is_extracted(name, filenames): + return name.lower().endswith('.zip') or not len(filenames) or is_in(name, filenames) + + def _extract_size(self, zip_paths, filenames): + total = 0 + for zip_path in zip_paths: + with ZipFile(zip_path, 'r') as z: + total += sum(i.file_size for i in z.infolist() if self._is_extracted(i.filename, filenames)) + return total + def _extract_zip(self, source_path, expand_path, filenames): ''' Extract matching entries, returning the subset of `filenames` (lower-cased) that were found by exact name - used by @@ -310,9 +325,10 @@ def _extract_zip(self, source_path, expand_path, filenames): # zip is itself extracted. if name.lower().endswith('.zip'): z.extract(zinfo, expand_path) + progress.reporter.advance(zinfo.file_size) continue - if len(filenames) and not is_in(name, filenames): + if not self._is_extracted(name, filenames): # Download only the named file, if any. _L.debug("Skipped file {}".format(name)) continue @@ -321,6 +337,7 @@ def _extract_zip(self, source_path, expand_path, filenames): found.add(name.lower()) z.extract(name, expand_path) + progress.reporter.advance(zinfo.file_size) return found def _find_single_zips(self, root_path): @@ -693,6 +710,10 @@ def ogr_source_to_csv(source_config, source_path, dest_path, disable_centroids=F writer = csv.DictWriter(f, fieldnames=out_fieldnames) writer.writeheader() + feature_count = in_layer.GetFeatureCount() + progress.reporter.start_phase('Reading source', total=feature_count if feature_count >= 0 else None) + row_count = 0 + in_feature = in_layer.GetNextFeature() while in_feature: row = dict() @@ -734,11 +755,15 @@ def ogr_source_to_csv(source_config, source_path, dest_path, disable_centroids=F row[GEOM_FIELDNAME] = None writer.writerow(row) + row_count += 1 + progress.reporter.advance() in_feature.Destroy() in_feature = in_layer.GetNextFeature() + progress.reporter.end_phase() in_datasource.Destroy() + return row_count def csv_source_to_csv(source_config, source_path, dest_path, disable_centroids=False): "Convert a source CSV file to an intermediate form, coerced to UTF-8 and EPSG:4326" @@ -808,6 +833,8 @@ def csv_source_to_csv(source_config, source_path, dest_path, disable_centroids=F with open(dest_path, 'w', encoding='utf-8') as dest_fp: writer = csv.DictWriter(dest_fp, out_fieldnames) writer.writeheader() + progress.reporter.start_phase('Reading source', unit='rows') + written = 0 # For every row in the source CSV row_number = 0 for source_row in reader: @@ -822,6 +849,11 @@ def csv_source_to_csv(source_config, source_path, dest_path, disable_centroids=F raise else: writer.writerow(out_row) + written += 1 + progress.reporter.advance() + + progress.reporter.end_phase() + return written def geojson_source_to_csv(source_config, source_path, dest_path, disable_centroids=False): ''' @@ -845,7 +877,10 @@ def geojson_source_to_csv(source_config, source_path, dest_path, disable_centroi with open(dest_path, 'w', encoding='utf-8') as dest_fp: writer = csv.DictWriter(dest_fp, out_fieldnames) writer.writeheader() + progress.reporter.start_phase('Reading source', total=os.path.getsize(source_path), unit='bytes') + written = 0 for (row_number, feature) in enumerate(stream_geojson(file)): + progress.reporter.set_done(file.buffer.tell()) try: row = feature['properties'] if feature['geometry'] is None: @@ -864,6 +899,10 @@ def geojson_source_to_csv(source_config, source_path, dest_path, disable_centroi else: row.update({GEOM_FIELDNAME: geom.ExportToWkt()}) writer.writerow(row) + written += 1 + + progress.reporter.end_phase() + return written _transform_cache = {} def _transform_to_4326(srs): @@ -1322,44 +1361,67 @@ def extract_to_source_csv(source_config, source_path, extract_path, disable_cent extract_path: file to write the extracted CSV file The extracted file will be in UTF-8 and will have X and Y columns corresponding - to longitude and latitude in EPSG:4326. + to longitude and latitude in EPSG:4326. Returns the number of rows written. """ format_string = source_config.data_source["conform"]['format'] protocol_string = source_config.data_source['protocol'] if format_string in ("shapefile", "xml", "gdb", "gpkg", "kml"): ogr_source_path = normalize_ogr_filename_case(source_path) - ogr_source_to_csv(source_config, ogr_source_path, extract_path, disable_centroids) + return ogr_source_to_csv(source_config, ogr_source_path, extract_path, disable_centroids) elif format_string == "csv": - csv_source_to_csv(source_config, source_path, extract_path, disable_centroids) + return csv_source_to_csv(source_config, source_path, extract_path, disable_centroids) elif format_string == "geojson": # GeoJSON sources have some awkward legacy with ESRI, see issue #34 if protocol_string == "ESRI": _L.info("ESRI GeoJSON source found; treating it as CSV") - csv_source_to_csv(source_config, source_path, extract_path, disable_centroids) + return csv_source_to_csv(source_config, source_path, extract_path, disable_centroids) else: _L.info("Non-ESRI GeoJSON source found; converting as a stream.") geojson_source_path = normalize_ogr_filename_case(source_path) - geojson_source_to_csv(source_config, geojson_source_path, extract_path, disable_centroids) + return geojson_source_to_csv(source_config, geojson_source_path, extract_path, disable_centroids) else: raise Exception("Unsupported source format %s" % format_string) -def transform_to_out_geojson(source_config, extract_path, dest_path): +def _report_geometry_cell(geometry): + if geometry: + shp = shape(geometry) + if not shp.is_empty: + point = shp if shp.geom_type == 'Point' else shp.representative_point() + progress.reporter.cell(point.x, point.y, 'ok') + +def _report_row_cell(wkt, outcome): + try: + point = wkt_loads(wkt) + progress.reporter.cell(point.x, point.y, outcome) + except Exception: + pass + +def transform_to_out_geojson(source_config, extract_path, dest_path, row_count=None): ''' Transform an extracted source CSV to the OpenAddresses output GeoJSON by applying conform rules. source_config: description of the source, containing the conform object extract_path: extracted CSV file to process dest_path: path for output file in OpenAddress CSV + row_count: number of rows in the extract, if known, for progress reporting ''' # Read through the extract CSV with open(extract_path, 'r', encoding='utf-8') as extract_fp: reader = csv.DictReader(extract_fp) # Write to the destination GeoJSON with open(dest_path, 'w', encoding='utf-8') as dest_fp: + progress.reporter.start_phase('Mapping columns', total=row_count, unit='rows') # For every row in the extract for extract_row in reader: - out_row = row_transform_and_convert(source_config, extract_row) + try: + out_row = row_transform_and_convert(source_config, extract_row) + except Exception: + _report_row_cell(extract_row.get(GEOM_FIELDNAME), 'error') + raise dest_fp.write(json.dumps(out_row, allow_nan=False) + '\n') + progress.reporter.advance() + _report_geometry_cell(out_row['geometry']) + progress.reporter.end_phase() def conform_cli(source_config, source_path, dest_path, disable_centroids=False): "Command line entry point for conforming a downloaded source to an output CSV." @@ -1380,8 +1442,8 @@ def conform_cli(source_config, source_path, dest_path, disable_centroids=False): _L.debug('extract temp file %s', extract_path) try: - extract_to_source_csv(source_config, source_path, extract_path, disable_centroids) - transform_to_out_geojson(source_config, extract_path, dest_path) + row_count = extract_to_source_csv(source_config, source_path, extract_path, disable_centroids) + transform_to_out_geojson(source_config, extract_path, dest_path, row_count) finally: os.remove(extract_path) diff --git a/openaddr/progress.py b/openaddr/progress.py new file mode 100644 index 00000000..c9fca939 --- /dev/null +++ b/openaddr/progress.py @@ -0,0 +1,101 @@ +import json +import math +import os +import time + +import h3 + +DEFAULT_INTERVAL = 15.0 +CELL_RESOLUTION = 6 +OUTCOME_INDEX = {'ok': 0, 'error': 2} + + +class Reporter: + ''' Prints machine-readable progress lines to stdout. + + Each line is a JSON object starting with {"progress": and is written + outside the logging module. An interval of 0 disables all output. + ''' + + def __init__(self, interval=None, clock=time.monotonic): + self._interval = interval + self._clock = clock + self._phase = None + + @property + def interval(self): + if self._interval is not None: + return self._interval + try: + return float(os.environ.get('OPENADDR_PROGRESS_INTERVAL', DEFAULT_INTERVAL)) + except ValueError: + return DEFAULT_INTERVAL + + def start_phase(self, name, total=None, unit='features', index=None, count=None): + interval = self.interval + if interval <= 0: + self._phase = None + return + self._phase = { + 'phase': name, 'phase_index': index, 'phase_count': count, + 'done': 0, 'total': total, 'unit': unit, + 'counters': {}, 'cells': {}, 'interval': interval, + } + self._emit() + + def advance(self, n=1): + if self._phase: + self._phase['done'] += n + self._emit_if_due() + + def add_total(self, n): + if self._phase: + self._phase['total'] = (self._phase['total'] or 0) + n + + def set_done(self, done): + if self._phase: + self._phase['done'] = done + self._emit_if_due() + + def count(self, counter_name, n=1): + if self._phase: + counters = self._phase['counters'] + counters[counter_name] = counters.get(counter_name, 0) + n + + def cell(self, lon, lat, outcome): + if not self._phase: + return + try: + if not (math.isfinite(lon) and math.isfinite(lat)): + return + cell_id = h3.latlng_to_cell(lat, lon, CELL_RESOLUTION) + except (TypeError, ValueError, h3.H3BaseException): + return + self._phase['cells'].setdefault(cell_id, [0, 0, 0])[OUTCOME_INDEX[outcome]] += 1 + + def end_phase(self): + if self._phase: + self._emit() + self._phase = None + + def _emit_if_due(self): + if self._clock() - self._phase['last_emit'] >= self._phase['interval']: + self._emit() + + def _emit(self): + phase = self._phase + report = {'phase': phase['phase']} + for key in ('phase_index', 'phase_count'): + if phase[key] is not None: + report[key] = phase[key] + report.update(done=phase['done'], total=phase['total'], unit=phase['unit']) + if phase['counters']: + report['counters'] = phase['counters'] + if phase['cells']: + report['cells'] = phase['cells'] + phase['cells'] = {} + phase['last_emit'] = self._clock() + print(json.dumps({'progress': report}, separators=(',', ':')), flush=True) + + +reporter = Reporter() diff --git a/openaddr/tests/__init__.py b/openaddr/tests/__init__.py index 56621043..36e51917 100644 --- a/openaddr/tests/__init__.py +++ b/openaddr/tests/__init__.py @@ -74,6 +74,8 @@ def return_path_in_second_arg_dir(_, path, *args, **kwargs): file.write('yo') return os.path.join(path, "out.geojson") +os.environ.setdefault('OPENADDR_PROGRESS_INTERVAL', '0') + class TestOA (unittest.TestCase): def setUp(self): diff --git a/openaddr/tests/cache.py b/openaddr/tests/cache.py index 5944a347..13991eea 100644 --- a/openaddr/tests/cache.py +++ b/openaddr/tests/cache.py @@ -415,6 +415,28 @@ def test_handle_feature_server_with_lat_lon_in_conform(self): self.assertTrue('oa:geom' in all_data[0]) self.assertEqual(all_data[0]['oa:geom'], 'POINT (-86.82960553 34.18671398)') + def test_esri_download_reports_progress(self): + """ ESRI Caching Emits Progress Lines With The Layer Feature Count As Total """ + import io + from contextlib import redirect_stdout + from .. import progress + + task = EsriRestDownloadTask('us-al-cullman') + c = SourceConfig(dict({ + "schema": 2, + "layers": {"addresses": [{"name": "default", "conform": {"lat": "LAT", "lon": "LON"}}]} + }), "addresses", "default") + out = io.StringIO() + with patch.object(progress, 'reporter', progress.Reporter(interval=15)), redirect_stdout(out): + with httmock.HTTMock(self.response_content): + task.download(["https://web2.kcsgis.com/kcsgis/rest/services/Cullman/VAM_Cullman_FS/FeatureServer/4"], self.workdir, c) + + lines = [json.loads(line)['progress'] for line in out.getvalue().splitlines()] + self.assertEqual(len(lines), 2) + self.assertEqual((lines[0]['phase'], lines[0]['done'], lines[0]['total']), ('Downloading from ArcGIS', 0, 5)) + self.assertEqual((lines[1]['done'], lines[1]['total']), (5, 5)) + self.assertEqual(sum(counts[0] for counts in lines[1]['cells'].values()), 5) + def test_skip_esri_features_with_nan_geometry(self): """ ESRI Caching Will Skip Features Whose Geometry Is The String "NaN" """ task = EsriRestDownloadTask('us-pa-bradford') diff --git a/openaddr/tests/progress.py b/openaddr/tests/progress.py new file mode 100644 index 00000000..6b53a437 --- /dev/null +++ b/openaddr/tests/progress.py @@ -0,0 +1,143 @@ +import io +import json +import csv +import os +import shutil +import tempfile +import unittest +from contextlib import redirect_stdout +from unittest.mock import patch + +from .. import SourceConfig, progress +from ..conform import GEOM_FIELDNAME, transform_to_out_geojson +from ..progress import Reporter + + +class FakeClock: + def __init__(self): + self.now = 0.0 + + def __call__(self): + return self.now + + +def run(reporter_fn, interval=15, clock=None): + reporter = Reporter(interval=interval, clock=clock or FakeClock()) + out = io.StringIO() + with redirect_stdout(out): + reporter_fn(reporter) + return [json.loads(line)['progress'] for line in out.getvalue().splitlines()], out.getvalue() + + +class TestProgressReporter(unittest.TestCase): + + def test_line_shape(self): + def script(r): + r.start_phase('Downloading', total=10, unit='bytes', index=1, count=3) + r.advance(4) + r.count('skipped', 2) + r.end_phase() + + lines, raw = run(script) + self.assertTrue(all(line.startswith('{"progress":') for line in raw.splitlines())) + self.assertNotIn(' ', raw.splitlines()[0]) + self.assertEqual(lines[0], {'phase': 'Downloading', 'phase_index': 1, 'phase_count': 3, + 'done': 0, 'total': 10, 'unit': 'bytes'}) + self.assertEqual(lines[-1]['done'], 4) + self.assertEqual(lines[-1]['counters'], {'skipped': 2}) + + def test_unknown_total_and_no_index(self): + lines, _ = run(lambda r: (r.start_phase('Reading source'), r.end_phase())) + self.assertIsNone(lines[0]['total']) + self.assertNotIn('phase_index', lines[0]) + self.assertEqual(lines[0]['unit'], 'features') + + def test_throttle(self): + clock = FakeClock() + + def script(r): + r.start_phase('p') + for _ in range(5): + clock.now += 4 + r.advance() + r.end_phase() + + lines, _ = run(script, clock=clock) + self.assertEqual([line['done'] for line in lines], [0, 4, 5]) + + def test_cell_deltas_reset_after_each_line(self): + clock = FakeClock() + + def script(r): + r.start_phase('p') + r.cell(-77.0, 38.9, 'ok') + r.cell(-77.0, 38.9, 'ok') + r.cell(-77.0, 38.9, 'error') + clock.now += 20 + r.advance() + r.cell(-77.0, 38.9, 'ok') + r.end_phase() + + lines, _ = run(script, clock=clock) + self.assertNotIn('cells', lines[0]) + (cell_id, counts), = lines[1]['cells'].items() + self.assertEqual(len(cell_id), 15) + self.assertEqual(counts, [2, 0, 1]) + self.assertEqual(lines[2]['cells'], {cell_id: [1, 0, 0]}) + + def test_disabled_when_interval_is_zero(self): + def script(r): + r.start_phase('p', total=1) + r.advance() + r.cell(1.0, 1.0, 'ok') + r.end_phase() + + _, raw = run(script, interval=0) + self.assertEqual(raw, '') + + def test_environment_controls_interval(self): + with patch.dict(os.environ, {'OPENADDR_PROGRESS_INTERVAL': '0'}): + self.assertEqual(Reporter().interval, 0) + with patch.dict(os.environ, {'OPENADDR_PROGRESS_INTERVAL': '3'}): + self.assertEqual(Reporter().interval, 3) + + def test_nonfinite_coordinates_are_ignored(self): + def script(r): + r.start_phase('p') + r.cell(float('nan'), 1.0, 'ok') + r.cell(1.0, float('inf'), 'error') + r.cell(None, 1.0, 'ok') + r.end_phase() + + lines, _ = run(script) + self.assertNotIn('cells', lines[-1]) + + +class TestProgressCallSites(unittest.TestCase): + + def test_transform_emits_first_and_final_line(self): + config = SourceConfig({ + 'schema': 2, + 'layers': {'addresses': [{'name': 'default', 'conform': {'street': 's', 'number': 'n'}}]}, + }, 'addresses', 'default') + workdir = tempfile.mkdtemp(prefix='testProgress-') + try: + extract_path = os.path.join(workdir, 'extract.csv') + with open(extract_path, 'w', encoding='utf-8') as file: + writer = csv.DictWriter(file, fieldnames=['s', 'n', GEOM_FIELDNAME]) + writer.writeheader() + for number in range(3): + writer.writerow({'s': 'DESMOND ST', 'n': str(number), GEOM_FIELDNAME: 'POINT (-76.51522 41.97925)'}) + + reporter = Reporter(interval=15) + out = io.StringIO() + with patch.object(progress, 'reporter', reporter), redirect_stdout(out): + transform_to_out_geojson(config, extract_path, os.path.join(workdir, 'out.geojson'), 3) + finally: + shutil.rmtree(workdir) + + lines = [json.loads(line)['progress'] for line in out.getvalue().splitlines()] + self.assertEqual(len(lines), 2) + self.assertEqual((lines[0]['phase'], lines[0]['done'], lines[0]['total']), ('Mapping columns', 0, 3)) + self.assertEqual((lines[1]['done'], lines[1]['total'], lines[1]['unit']), (3, 3, 'rows')) + self.assertEqual(list(lines[1]['cells'].values()), [[3, 0, 0]]) diff --git a/setup.py b/setup.py index 5620e8e1..e214fc66 100644 --- a/setup.py +++ b/setup.py @@ -55,6 +55,8 @@ # https://github.com/tilezen/mapbox-vector-tile 'mapbox-vector-tile == 2.0.1', + # H3 cell ids for progress reporting + 'h3 >= 4, < 5', 'future==0.18.3', 'protobuf>=4.21, <5', ] diff --git a/test.py b/test.py index 5fb2b7a3..ca7f127d 100755 --- a/test.py +++ b/test.py @@ -17,6 +17,7 @@ from openaddr.tests.cache import TestCacheExtensionGuessing, TestCacheEsriDownload from openaddr.tests.conform import TestConformCli, TestConformTransforms, TestConformMisc, TestConformCsv, TestConformTests from openaddr.tests.preview import TestPreview +from openaddr.tests.progress import TestProgressReporter, TestProgressCallSites from openaddr.tests.slippymap import TestSlippyMap from openaddr.tests.util import TestUtilities