Repository: symerio/pgeocode

Issue:
Support alternate download locations
It might be useful to support alternate download locations in case GeoNames website goes down. This would also help reproducibility (I'm not sure how often GeoNames database is updated and if that is tracked somewhere).

This would require storing the data somewhere. One possibility for free hosting could be to attach it to Github releases.

For instance, maybe @zaro's implementation in https://github.com/zaro/pgeocode/commit/6a3c743bee8fd67ae6ec82c87e0d6cbfefa62110 could be adapted.

Comments:

Patch:
diff --git a/README.rst b/README.rst
index b1b279c..a3af1ab 100644
--- a/README.rst
+++ b/README.rst
@@ -103,6 +103,27 @@ Defaults to ``~/pgeocode_data``, it is the directory where data is downloaded
 for later consumption. It can be changed using the environment variable
 ``PGEOCODE_DATA_DIR``, i.e. ``export PGEOCODE_DATA_DIR=/tmp/pgeocode_data``.
 
+**Data sources**
+
+The data sources are provided as a list in the ``pgeocode.DOWNLOAD_URL`` variable.
+The default value is,
+
+.. code::
+
+    DOWNLOAD_URL = [
+        "https://download.geonames.org/export/zip/{country}.zip",
+        "https://symerio.github.io/postal-codes-data/data/geonames/{country}.txt",
+    ]
+
+Data sources are tried from first to last until one works. Here the second link is a mirror
+of the first.
+
+It is also possible to extend this variable with third party data sources, as
+long as they follow the same format. See for instance
+[postal-codes-data](https://github.com/symerio/postal-codes-data/tree/master/data/geonames)
+repository for examples of data files.
+
+
 License
 -------
 
diff --git a/doc/contributing.rst b/doc/contributing.rst
index 914ab51..35f3408 100644
--- a/doc/contributing.rst
+++ b/doc/contributing.rst
@@ -9,6 +9,10 @@ Testing
 
 Unit tests can be run with,
 
+.. code::
+
+    pip install pytest pytest-httpserver
+
 .. code::
 
     pytest
diff --git a/pgeocode.py b/pgeocode.py
index 191a37d..fe6a6d9 100644
--- a/pgeocode.py
+++ b/pgeocode.py
@@ -2,11 +2,13 @@
 #
 # Authors: Roman Yurchak <roman.yurchak@symerio.com>
 
+import contextlib
 import os
 import urllib.request
 import warnings
 from io import BytesIO
-from typing import Any, Tuple
+from typing import Any, Tuple, List
+from zipfile import ZipFile
 
 import numpy as np
 import pandas as pd
@@ -17,7 +19,13 @@ STORAGE_DIR = os.environ.get(
     "PGEOCODE_DATA_DIR", os.path.join(os.path.expanduser("~"), "pgeocode_data")
 )
 
-DOWNLOAD_URL = "https://download.geonames.org/export/zip/{country}.zip"
+# A list of download locations. If the first URL fails, following ones will
+# be used.
+DOWNLOAD_URL = [
+    "https://download.geonames.org/export/zip/{country}.zip",
+    "https://symerio.github.io/postal-codes-data/data/geonames/{country}.txt",
+]
+
 
 DATA_FIELDS = [
     "country_code",
@@ -121,11 +129,51 @@ COUNTRIES_VALID = [
 ]
 
 
-def _open_url(url: str) -> Tuple[BytesIO, Any]:
-    """Download contents for a URL"""
+@contextlib.contextmanager
+def _open_extract_url(url: str, country: str) -> Any:
+    """Download contents for a URL
+
+    If the file has a .zip extension, open it and extract the country
+
+    Returns the opened file object.
+    """
     with urllib.request.urlopen(url) as res:
-        reader = BytesIO(res.read())
-    return reader, res.headers
+        with BytesIO(res.read()) as reader:
+            if url.endswith(".zip"):
+                with ZipFile(reader) as fh_zip:
+                    with fh_zip.open(country.upper() + ".txt") as fh:
+                        yield fh
+            else:
+                yield reader
+
+
+@contextlib.contextmanager
+def _open_extract_cycle_url(urls: List[str], country: str) -> Any:
+    """Same as _open_extract_url but cycle through URLs until one works
+
+    We start by opening the first URL in the list, and if fails
+    move to the next, until one works or the end of list is reached.
+    """
+    if not isinstance(urls, list) or not len(urls):
+        raise ValueError(f"urls={urls} must be a list with at least one URL")
+
+    err_msg = f"Provided download URLs failed {{err}}: {urls}"
+    for idx, val in enumerate(urls):
+        try:
+            with _open_extract_url(val, country) as fh:
+                yield fh
+            # Found a working URL, exit the loop.
+            break
+        except urllib.error.HTTPError as err:  # type: ignore
+            if idx == len(urls) - 1:
+                raise
+            warnings.warn(
+                f"Download from {val} failed with: {err}. "
+                "Trying next URL in DOWNLOAD_URL list.",
+                UserWarning,
+            )
+    else:
+        raise ValueError(err_msg)
 
 
 class Nominatim:
@@ -168,23 +216,22 @@ class Nominatim:
     @staticmethod
     def _get_data(country: str) -> Tuple[str, pd.DataFrame]:
         """Load the data from disk; otherwise download and save it"""
-        from zipfile import ZipFile
 
         data_path = os.path.join(STORAGE_DIR, country.upper() + ".txt")
         if os.path.exists(data_path):
             data = pd.read_csv(data_path, dtype={"postal_code": str})
         else:
-            url = DOWNLOAD_URL.format(country=country)
-            reader, headers = _open_url(url)
-            with ZipFile(reader) as fh_zip:
-                with fh_zip.open(country.upper() + ".txt") as fh:
-                    data = pd.read_csv(
-                        fh,
-                        sep="\t",
-                        header=None,
-                        names=DATA_FIELDS,
-                        dtype={"postal_code": str},
-                    )
+            download_urls = [
+                val.format(country=country) for val in DOWNLOAD_URL
+            ]
+            with _open_extract_cycle_url(download_urls, country) as fh:
+                data = pd.read_csv(
+                    fh,
+                    sep="\t",
+                    header=None,
+                    names=DATA_FIELDS,
+                    dtype={"postal_code": str},
+                )
             if not os.path.exists(STORAGE_DIR):
                 os.mkdir(STORAGE_DIR)
             data.to_csv(data_path, index=None)


Test patch:
diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
index c52c3b3..d7259f8 100644
--- a/.github/workflows/python-tests.yml
+++ b/.github/workflows/python-tests.yml
@@ -38,5 +38,5 @@ jobs:
         pip install flake8
     - name: Test with pytest
       run: |
-        pip install pytest
+        pip install pytest pytest-httpserver
         pytest
diff --git a/test_pgeocode.py b/test_pgeocode.py
index d8c966a..fe64c6a 100644
--- a/test_pgeocode.py
+++ b/test_pgeocode.py
@@ -2,8 +2,10 @@
 #
 # Authors: Roman Yurchak <roman.yurchak@symerio.com>
 import os
-import shutil
-import tempfile
+import urllib
+import json
+from zipfile import ZipFile
+from io import BytesIO
 
 import numpy as np
 import pandas as pd
@@ -12,16 +14,13 @@ from numpy.testing import assert_allclose, assert_array_equal
 
 import pgeocode
 from pgeocode import GeoDistance, Nominatim, haversine_distance
+from pgeocode import _open_extract_url
 
 
 @pytest.fixture
-def temp_dir():
-    path_save = pgeocode.STORAGE_DIR
-    path = tempfile.mkdtemp()
-    pgeocode.STORAGE_DIR = path
-    yield path
-    pgeocode.STORAGE_DIR = path_save
-    shutil.rmtree(path)
+def temp_dir(tmpdir, monkeypatch):
+    monkeypatch.setattr(pgeocode, "STORAGE_DIR", str(tmpdir))
+    yield str(tmpdir)
 
 
 def _normalize_str(x):
@@ -179,3 +178,82 @@ def test_haversine_distance():
     d_pred = haversine_distance(x, y)
     # same distance +/- 3 km
     assert_allclose(d_ref, d_pred, atol=3)
+
+
+def test_open_extract_url(httpserver):
+    download_url = "/fr.txt"
+
+    # check download of uncompressed files
+    httpserver.expect_oneshot_request(download_url).respond_with_json({"a": 1})
+    with _open_extract_url(httpserver.url_for(download_url), "fr") as fh:
+        assert json.loads(fh.read()) == {"a": 1}
+    httpserver.check_assertions()
+
+    # check download of zipped files
+    # Create an in-memory zip file
+    answer = b"a=1"
+    with BytesIO() as fh:
+        with ZipFile(fh, "w") as fh_zip:
+            with fh_zip.open("FR.txt", "w") as fh_inner:
+                fh_inner.write(answer)
+        fh.seek(0)
+        res = fh.read()
+
+    download_url = "/fr.zip"
+    httpserver.expect_oneshot_request(download_url).respond_with_data(res)
+
+    with _open_extract_url(httpserver.url_for(download_url), "fr") as fh:
+        assert fh.read() == answer
+
+
+@pytest.mark.parametrize(
+    "download_url",
+    [
+        "https://download.geonames.org/export/zip/{country}.zip",
+        "https://symerio.github.io/postal-codes-data/data/"
+        "geonames/{country}.txt",
+    ],
+    ids=["geonames", "gitlab-pages"],
+)
+def test_cdn(temp_dir, monkeypatch, download_url):
+    monkeypatch.setattr(pgeocode, "DOWNLOAD_URL", [download_url])
+    assert not os.path.exists(os.path.join(temp_dir, "IE.txt"))
+    Nominatim("IE")
+    # the data file was downloaded
+    assert os.path.exists(os.path.join(temp_dir, "IE.txt"))
+
+
+def test_url_returns_404(httpserver, monkeypatch, temp_dir):
+    download_url = "/fr.gzip"
+    httpserver.expect_oneshot_request(download_url).respond_with_data(
+        "", status=404
+    )
+
+    monkeypatch.setattr(
+        pgeocode, "DOWNLOAD_URL", [httpserver.url_for(download_url)]
+    )
+    # Nominatim("fr")
+    with pytest.raises(urllib.error.HTTPError, match="HTTP Error 404"):
+        Nominatim("fr")
+    httpserver.check_assertions()
+
+
+def test_first_url_fails(httpserver, monkeypatch, temp_dir):
+    download_url = "/IE.txt"
+    httpserver.expect_oneshot_request(download_url).respond_with_data(
+        "", status=404
+    )
+
+    monkeypatch.setattr(
+        pgeocode,
+        "DOWNLOAD_URL",
+        [
+            httpserver.url_for(download_url),
+            "https://symerio.github.io/postal-codes-data/data/"
+            "geonames/{country}.txt",
+        ],
+    )
+    msg = "IE.txt failed with: HTTP Error 404.*Trying next URL"
+    with pytest.warns(UserWarning, match=msg):
+        Nominatim("ie")
+    httpserver.check_assertions()