geobatch 0.1.0__py3-none-any.whl

geobatch/__init__.py

@@ -0,0 +1,97 @@
+"""geobatch: polite, resumable batch geocoding and reverse geocoding.
+
+Quick start::
+
+    from geobatch import Geocoder, Nominatim
+
+    gc = Geocoder(Nominatim(user_agent="MyProject/1.0 (https://my-project.org)"))
+    result = gc.geocode_many(["Brandenburger Tor, Berlin", "Eiffel Tower, Paris"])
+    for rec in result:
+        print(rec.status, rec.result.lat if rec.result else None, rec.flags)
+"""
+
+from ._version import __version__
+from .batch import BatchInterrupted, BatchResult, Geocoder, Summary, to_query
+from .cache import Cache, default_cache_path
+from .errors import (
+    ConfigurationError,
+    GeobatchError,
+    HttpError,
+    MalformedResponseError,
+    MissingAPIKeyError,
+    ProviderError,
+    TransportError,
+    UsagePolicyError,
+)
+from .http import RetryPolicy
+from .io import (
+    build_queries,
+    read_table,
+    to_geojson,
+    write_csv,
+    write_geojson,
+    write_gpkg,
+    write_output,
+)
+from .models import BatchRecord, GeocodeResult, Query
+from .normalize import normalize_address, normalize_query
+from .providers import (
+    PROVIDERS,
+    Geoapify,
+    GoogleGeocoding,
+    Mapbox,
+    Nominatim,
+    OpenCage,
+    Pelias,
+    Photon,
+    Provider,
+    SearchOptions,
+    TemplateProvider,
+    make_provider,
+)
+from .quality import QualityConfig
+
+__all__ = [
+    "PROVIDERS",
+    "BatchInterrupted",
+    "BatchRecord",
+    "BatchResult",
+    "Cache",
+    "ConfigurationError",
+    "GeobatchError",
+    "GeocodeResult",
+    "Geocoder",
+    "Geoapify",
+    "GoogleGeocoding",
+    "HttpError",
+    "MalformedResponseError",
+    "Mapbox",
+    "MissingAPIKeyError",
+    "Nominatim",
+    "OpenCage",
+    "Pelias",
+    "Photon",
+    "Provider",
+    "ProviderError",
+    "QualityConfig",
+    "Query",
+    "RetryPolicy",
+    "SearchOptions",
+    "Summary",
+    "TemplateProvider",
+    "TransportError",
+    "UsagePolicyError",
+    "__version__",
+    "build_queries",
+    "default_cache_path",
+    "make_provider",
+    "normalize_address",
+    "normalize_query",
+    "read_table",
+    "to_geojson",
+    "to_query",
+    "write_csv",
+    "write_geojson",
+    "write_gpkg",
+    "write_output",
+]

geobatch/__main__.py

@@ -0,0 +1,7 @@
+"""``python -m geobatch``."""
+
+import sys
+
+from .cli import main
+
+sys.exit(main())

geobatch/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

geobatch/batch.py

@@ -0,0 +1,526 @@
+"""The batch orchestrator: cache, de-duplication, concurrency, fallback, flags."""
+
+from __future__ import annotations
+
+import os
+import sys
+import threading
+import time
+from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    TextIO,
+    Tuple,
+    Union,
+)
+
+from .cache import Cache, make_key
+from .errors import ConfigurationError, GeobatchError, ProviderError, UsagePolicyError
+from .models import (
+    STATUS_ERROR,
+    STATUS_NOT_FOUND,
+    STATUS_OK,
+    STATUS_PENDING,
+    STATUS_SKIPPED,
+    BatchRecord,
+    GeocodeResult,
+    Query,
+)
+from .normalize import normalize_query, query_cache_text
+from .progress import ProgressBar
+from .providers.base import Provider, SearchOptions
+from .quality import QualityConfig, assess_record, count_flags, flag_duplicates
+
+QueryLike = Union[str, Query, Mapping[str, Any], None]
+CacheLike = Union[Cache, str, "os.PathLike[str]", bool, None]
+
+
+class BatchInterrupted(KeyboardInterrupt):
+    """Raised on Ctrl+C during a batch; ``.result`` holds the partial result.
+
+    Every lookup finished before the interruption is already in the cache, so
+    running the same batch again resumes where it stopped.
+    """
+
+    def __init__(self, result: BatchResult) -> None:
+        super().__init__("batch interrupted")
+        self.result = result
+
+
+@dataclass
+class Summary:
+    """Counts describing one batch run."""
+
+    kind: str = "forward"
+    total: int = 0
+    ok: int = 0
+    not_found: int = 0
+    error: int = 0
+    skipped: int = 0
+    pending: int = 0
+    unique_lookups: int = 0
+    cache_hits: int = 0
+    network_requests: int = 0
+    retries: int = 0
+    by_provider: Dict[str, int] = field(default_factory=dict)
+    by_precision: Dict[str, int] = field(default_factory=dict)
+    flags: Dict[str, int] = field(default_factory=dict)
+    elapsed_s: float = 0.0
+    interrupted: bool = False
+    aborted: str = ""
+
+    def to_dict(self) -> Dict[str, Any]:
+        return dict(self.__dict__)
+
+    def format(self) -> str:
+        lines = [
+            f"geobatch {self.kind} summary",
+            f"  rows            {self.total}",
+            f"  ok              {self.ok}",
+            f"  not found       {self.not_found}",
+            f"  errors          {self.error}",
+            f"  skipped (empty) {self.skipped}",
+        ]
+        if self.pending:
+            lines.append(f"  not processed   {self.pending}")
+        lines += [
+            f"  unique lookups  {self.unique_lookups}",
+            f"  cache hits      {self.cache_hits}",
+            f"  HTTP requests   {self.network_requests} (retries {self.retries})",
+            f"  elapsed         {self.elapsed_s:.1f} s",
+        ]
+        if self.by_provider:
+            lines.append(
+                "  by provider     "
+                + ", ".join(f"{k}={v}" for k, v in sorted(self.by_provider.items()))
+            )
+        if self.by_precision:
+            lines.append(
+                "  by precision    "
+                + ", ".join(f"{k}={v}" for k, v in sorted(self.by_precision.items()))
+            )
+        if self.flags:
+            lines.append(
+                "  quality flags   " + ", ".join(f"{k}={v}" for k, v in sorted(self.flags.items()))
+            )
+        if self.interrupted:
+            lines.append("  INTERRUPTED: re-run the same command to resume from the cache")
+        if self.aborted:
+            lines.append(f"  ABORTED: {self.aborted}")
+        return "\n".join(lines)
+
+
+@dataclass
+class BatchResult:
+    records: List[BatchRecord]
+    summary: Summary
+
+    def __iter__(self) -> Any:
+        return iter(self.records)
+
+    def __len__(self) -> int:
+        return len(self.records)
+
+
+@dataclass
+class _Outcome:
+    status: str
+    results: List[GeocodeResult] = field(default_factory=list)
+    provider: str = ""
+    from_cache: bool = False
+    attempts: int = 0
+    error: str = ""
+
+
+def to_query(value: QueryLike) -> Query:
+    """Coerce a string, mapping or :class:`Query` into a :class:`Query`."""
+    if value is None:
+        return Query()
+    if isinstance(value, Query):
+        return value
+    if isinstance(value, str):
+        return Query(text=value)
+    if isinstance(value, Mapping):
+        allowed = {"text", "street", "city", "county", "state", "postcode", "country"}
+        unknown = set(value) - allowed
+        if unknown:
+            raise ConfigurationError(f"unknown query fields: {', '.join(sorted(unknown))}")
+        return Query(**{k: "" if v is None else str(v) for k, v in value.items()})
+    raise ConfigurationError(f"cannot turn {type(value).__name__} into a Query")
+
+
+def _open_cache(cache: CacheLike) -> Optional[Cache]:
+    if cache is None or cache is False:
+        return None
+    if cache is True:
+        return Cache()
+    if isinstance(cache, Cache):
+        return cache
+    return Cache(Path(cache))
+
+
+class Geocoder:
+    """Batch geocoder over a fallback chain of providers.
+
+    Parameters
+    ----------
+    providers:
+        One provider or a list tried in order (fallback chain).
+    cache:
+        ``True`` (default) for the per-user SQLite cache, a path, a
+        :class:`Cache`, or ``False`` to disable caching (refused for providers
+        whose policy requires caching, such as public Nominatim).
+    options:
+        Country/bbox biasing and language, see :class:`SearchOptions`.
+    quality:
+        Thresholds for quality flags; by default the expected countries and
+        bbox are taken from ``options``.
+    fallback_on_not_found:
+        Also ask the next provider when one finds nothing (not only on errors).
+    max_consecutive_errors:
+        Abort the batch after this many consecutive failed lookups (the
+        service is probably down or blocking you; hammering it is impolite).
+    notice_stream:
+        Where policy notices are printed (stderr by default, ``None`` = silent).
+    """
+
+    def __init__(
+        self,
+        providers: Union[Provider, Sequence[Provider]],
+        cache: CacheLike = True,
+        options: Optional[SearchOptions] = None,
+        quality: Optional[QualityConfig] = None,
+        fallback_on_not_found: bool = True,
+        max_consecutive_errors: int = 25,
+        notice_stream: Optional[TextIO] = sys.stderr,
+    ) -> None:
+        chain = [providers] if isinstance(providers, Provider) else list(providers)
+        if not chain:
+            raise ConfigurationError("at least one provider is required")
+        labels = [p.label for p in chain]
+        if len(set(labels)) != len(labels):
+            raise ConfigurationError("providers in a chain need distinct labels")
+        self.providers: List[Provider] = chain
+        self.cache = _open_cache(cache)
+        for p in chain:
+            if p.requires_cache and self.cache is None:
+                raise UsagePolicyError(
+                    f"{p.label}: results must be cached under this provider's usage policy; "
+                    "enable the cache"
+                )
+        self.options = options or SearchOptions()
+        self.quality = quality or QualityConfig(
+            expected_countries=self.options.country_codes,
+            expected_bbox=self.options.bbox if self.options.bounded else None,
+        )
+        self.fallback_on_not_found = fallback_on_not_found
+        self.max_consecutive_errors = max(1, int(max_consecutive_errors))
+        self.notices: List[str] = [n for p in chain for n in p.notices]
+        if notice_stream is not None:
+            for note in self.notices:
+                notice_stream.write(f"geobatch: {note}\n")
+            notice_stream.flush()
+
+    # ------------------------------------------------------------------ core
+    def _lookup(
+        self,
+        kind: str,
+        key_text: str,
+        call: Callable[[Provider], Tuple[List[GeocodeResult], int]],
+    ) -> _Outcome:
+        errors: List[str] = []
+        not_found_by: List[str] = []
+        attempts = 0
+        all_cached = True
+        for prov in self.providers:
+            key = make_key(prov.provider_id, kind, key_text, prov.cache_params(self.options))
+            use_cache = self.cache is not None and prov.cacheable
+            if use_cache:
+                entry = self.cache.get(key)  # type: ignore[union-attr]
+                if entry is not None:
+                    if entry.found:
+                        return _Outcome(STATUS_OK, entry.results, prov.label, all_cached, attempts)
+                    not_found_by.append(prov.label)
+                    if not self.fallback_on_not_found:
+                        break
+                    continue
+            all_cached = False
+            try:
+                results, n = call(prov)
+                attempts += n
+            except UsagePolicyError as exc:
+                errors.append(f"{prov.label}: {exc}")
+                continue
+            except ProviderError as exc:
+                errors.append(f"{prov.label}: {exc}")
+                continue
+            if use_cache:
+                self.cache.put(  # type: ignore[union-attr]
+                    key, prov.provider_id, kind, key_text, results, prov.cache_ttl
+                )
+            if results:
+                return _Outcome(STATUS_OK, results, prov.label, False, attempts)
+            not_found_by.append(prov.label)
+            if not self.fallback_on_not_found:
+                break
+        if not_found_by:
+            return _Outcome(
+                STATUS_NOT_FOUND,
+                provider=not_found_by[-1],
+                from_cache=all_cached,
+                attempts=attempts,
+                error="; ".join(errors),
+            )
+        return _Outcome(STATUS_ERROR, attempts=attempts, error="; ".join(errors))
+
+    def _run(
+        self,
+        kind: str,
+        records: List[BatchRecord],
+        jobs: Dict[str, Tuple[Callable[[Provider], Tuple[List[GeocodeResult], int]], List[int]]],
+        progress: Union[bool, ProgressBar],
+        expected: Optional[Sequence[str]],
+    ) -> BatchResult:
+        start = time.monotonic()
+        before = [(p.transport.stats.requests, p.transport.stats.retries) for p in self.providers]
+        summary = Summary(kind=kind, total=len(records), unique_lookups=len(jobs))
+        bar: Optional[ProgressBar]
+        if isinstance(progress, ProgressBar):
+            bar = progress
+        elif progress:
+            bar = ProgressBar(len(records), label="geocoding" if kind == "forward" else "reverse")
+        else:
+            bar = None
+        if bar is not None:
+            skipped = sum(1 for r in records if r.status == STATUS_SKIPPED)
+            if skipped:
+                bar.update(skipped)
+
+        abort = threading.Event()
+        abort_reason: List[str] = []
+        streak = [0]
+        streak_lock = threading.Lock()
+
+        def task(key_text: str) -> _Outcome:
+            if abort.is_set():
+                return _Outcome(STATUS_PENDING, error=abort_reason[0] if abort_reason else "")
+            call = jobs[key_text][0]
+            outcome = self._lookup(kind, key_text, call)
+            with streak_lock:
+                if outcome.status == STATUS_ERROR:
+                    streak[0] += 1
+                    if streak[0] >= self.max_consecutive_errors and not abort.is_set():
+                        abort_reason.append(
+                            f"{streak[0]} consecutive lookups failed (last: {outcome.error}); "
+                            "stopping to avoid hammering the service"
+                        )
+                        abort.set()
+                else:
+                    streak[0] = 0
+            return outcome
+
+        def apply(key_text: str, outcome: _Outcome) -> None:
+            for idx in jobs[key_text][1]:
+                rec = records[idx]
+                rec.status = outcome.status
+                rec.result = outcome.results[0] if outcome.results else None
+                rec.provider = outcome.provider
+                rec.from_cache = outcome.from_cache and outcome.status != STATUS_PENDING
+                rec.attempts = outcome.attempts
+                rec.error = outcome.error
+                if bar is not None and outcome.status != STATUS_PENDING:
+                    bar.update(cached=rec.from_cache, failed=outcome.status == STATUS_ERROR)
+
+        workers = max(p.concurrency for p in self.providers)
+        pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="geobatch")
+        pending: Dict[Future[_Outcome], str] = {}
+        interrupted = False
+        try:
+            keys = iter(list(jobs))
+            # Keep a bounded number of tasks in flight so Ctrl+C stops quickly.
+            for key_text in keys:
+                pending[pool.submit(task, key_text)] = key_text
+                if len(pending) >= workers * 2:
+                    break
+            while pending:
+                done, _ = wait(list(pending), return_when=FIRST_COMPLETED)
+                for fut in done:
+                    apply(pending.pop(fut), fut.result())
+                    nxt = next(keys, None)
+                    if nxt is not None:
+                        pending[pool.submit(task, nxt)] = nxt
+        except KeyboardInterrupt:
+            interrupted = True
+            abort_reason.append("interrupted")
+            abort.set()
+            for fut in pending:
+                fut.cancel()
+        finally:
+            # Requests already on the wire are allowed to finish (and be
+            # cached); queued work is cancelled or returns immediately.
+            pool.shutdown(wait=True, cancel_futures=True)
+            if bar is not None:
+                bar.close()
+
+        self._finish(records, summary, expected)
+        summary.elapsed_s = round(time.monotonic() - start, 3)
+        summary.network_requests = sum(
+            p.transport.stats.requests - b[0] for p, b in zip(self.providers, before)
+        )
+        summary.retries = sum(
+            p.transport.stats.retries - b[1] for p, b in zip(self.providers, before)
+        )
+        summary.cache_hits = len(
+            {query_key for query_key, (_, idxs) in jobs.items() if records[idxs[0]].from_cache}
+        )
+        summary.interrupted = interrupted
+        if abort_reason and not interrupted:
+            summary.aborted = abort_reason[0]
+        result = BatchResult(records, summary)
+        if interrupted:
+            raise BatchInterrupted(result)
+        return result
+
+    def _finish(
+        self,
+        records: List[BatchRecord],
+        summary: Summary,
+        expected: Optional[Sequence[str]],
+    ) -> None:
+        first = self.providers[0].label
+        for rec in records:
+            exp = expected[rec.index] if expected is not None and rec.index < len(expected) else ""
+            assess_record(rec, self.quality, expected_country=exp or "", first_provider=first)
+        flag_duplicates(records, self.quality.duplicate_decimals)
+        for rec in records:
+            if rec.status == STATUS_OK:
+                summary.ok += 1
+                summary.by_provider[rec.provider] = summary.by_provider.get(rec.provider, 0) + 1
+                if rec.result is not None:
+                    p = rec.result.precision
+                    summary.by_precision[p] = summary.by_precision.get(p, 0) + 1
+            elif rec.status == STATUS_NOT_FOUND:
+                summary.not_found += 1
+            elif rec.status == STATUS_ERROR:
+                summary.error += 1
+            elif rec.status == STATUS_SKIPPED:
+                summary.skipped += 1
+            else:
+                summary.pending += 1
+        summary.flags = count_flags(records)
+
+    # ------------------------------------------------------------------ public
+    def geocode_many(
+        self,
+        queries: Iterable[QueryLike],
+        rows: Optional[Sequence[Dict[str, Any]]] = None,
+        expected_countries: Optional[Sequence[str]] = None,
+        progress: Union[bool, ProgressBar] = False,
+    ) -> BatchResult:
+        """Forward-geocode many queries; identical queries are looked up once.
+
+        ``rows`` (optional) are attached to the records for output writers;
+        ``expected_countries`` gives a per-row ISO code used for the
+        ``country_mismatch`` flag.
+        """
+        records: List[BatchRecord] = []
+        jobs: Dict[
+            str, Tuple[Callable[[Provider], Tuple[List[GeocodeResult], int]], List[int]]
+        ] = {}
+        for i, raw in enumerate(queries):
+            query = normalize_query(to_query(raw))
+            rec = BatchRecord(index=i, query=query, row=dict(rows[i]) if rows else {})
+            records.append(rec)
+            if query.is_empty:
+                rec.status = STATUS_SKIPPED
+                rec.error = "empty address"
+                continue
+            key_text = query_cache_text(query)
+            if key_text in jobs:
+                jobs[key_text][1].append(i)
+            else:
+                jobs[key_text] = (self._forward_call(query), [i])
+        return self._run("forward", records, jobs, progress, expected_countries)
+
+    def _forward_call(self, query: Query) -> Callable[[Provider], Tuple[List[GeocodeResult], int]]:
+        options = self.options
+        return lambda prov: prov.geocode(query, options)
+
+    def _reverse_call(
+        self, lat: float, lon: float
+    ) -> Callable[[Provider], Tuple[List[GeocodeResult], int]]:
+        options = self.options
+        return lambda prov: prov.reverse(lat, lon, options)
+
+    def reverse_many(
+        self,
+        points: Iterable[Tuple[Any, Any]],
+        rows: Optional[Sequence[Dict[str, Any]]] = None,
+        progress: Union[bool, ProgressBar] = False,
+        decimals: int = 6,
+    ) -> BatchResult:
+        """Reverse-geocode ``(lat, lon)`` pairs; invalid pairs are skipped.
+
+        Points are rounded to ``decimals`` places (6 ~ 0.1 m) for caching and
+        de-duplication.
+        """
+        records: List[BatchRecord] = []
+        jobs: Dict[
+            str, Tuple[Callable[[Provider], Tuple[List[GeocodeResult], int]], List[int]]
+        ] = {}
+        for i, (lat_raw, lon_raw) in enumerate(points):
+            rec = BatchRecord(index=i, row=dict(rows[i]) if rows else {})
+            records.append(rec)
+            try:
+                lat, lon = float(lat_raw), float(lon_raw)
+            except (TypeError, ValueError):
+                lat = lon = float("nan")
+            if not (-90 <= lat <= 90 and -180 <= lon <= 180):
+                rec.status = STATUS_SKIPPED
+                rec.error = "invalid coordinates"
+                continue
+            lat, lon = round(lat, decimals), round(lon, decimals)
+            rec.point = (lat, lon)
+            key_text = f"{lat:.{decimals}f},{lon:.{decimals}f}"
+            if key_text in jobs:
+                jobs[key_text][1].append(i)
+            else:
+                jobs[key_text] = (self._reverse_call(lat, lon), [i])
+        return self._run("reverse", records, jobs, progress, None)
+
+    def geocode(self, query: QueryLike) -> BatchRecord:
+        """Geocode one address (with cache, fallback and flags)."""
+        return self.geocode_many([query]).records[0]
+
+    def reverse(self, lat: float, lon: float) -> BatchRecord:
+        """Reverse-geocode one point (with cache, fallback and flags)."""
+        return self.reverse_many([(lat, lon)]).records[0]
+
+    def close(self) -> None:
+        if self.cache is not None:
+            self.cache.close()
+
+    def __enter__(self) -> Geocoder:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+
+__all__ = [
+    "BatchInterrupted",
+    "BatchResult",
+    "GeobatchError",
+    "Geocoder",
+    "Summary",
+    "to_query",
+]