From 8e25e76c82880b111df5a3b3b4b4daa8909694ba Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Simon=20H=C3=B8xbro=20Hansen?= <simon.hansen@me.com>
Date: Wed, 14 Aug 2024 11:54:11 +0200
Subject: [PATCH] Fix download data for examples (#155)

---
 .github/workflows/test.yaml |  2 ++
 examples/download.py        | 34 +++++++++++++++++++++++++++-------
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 3a2bc7e..a7ce028 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -130,6 +130,8 @@ jobs:
           pixi run -e ${{ matrix.environment }} test-unit $COV
       - name: Test Examples
         if: needs.setup.outputs.code_change == 'true'
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}  # For download data
         run: |
           pixi run -e ${{ matrix.environment }} test-example
       - uses: codecov/codecov-action@v4
diff --git a/examples/download.py b/examples/download.py
index 4715cdc..a217d8a 100644
--- a/examples/download.py
+++ b/examples/download.py
@@ -1,25 +1,37 @@
+import os
+from shutil import rmtree
+
 try:
     import requests
     from platformdirs import user_cache_path
 except ImportError:
-    raise ImportError("requests and platformdirs are needed to download data")
+    raise ImportError("requests and platformdirs are needed to download data") from None
+
+
+if os.environ.get("GITHUB_TOKEN"):
+    HEADERS = {"Authorization": f"token {os.environ['GITHUB_TOKEN']}"}
+else:
+    HEADERS = None
 
 
 def download_map(dataset):
     if dataset not in ("naturalearth_lowres", "naturalearth_cities"):
-        raise ValueError(f"Unknown dataset: {dataset}, supported datasets are 'naturalearth_lowres' and 'naturalearth_cities'")
+        raise ValueError(
+            f"Unknown dataset: {dataset}, supported datasets are 'naturalearth_lowres' and 'naturalearth_cities'"
+        )
     url = f"https://api.github.com/repos/geopandas/geopandas/contents/geopandas/datasets/{dataset}?ref=v0.14.4"
     local_dir = user_cache_path() / "spatialpandas" / dataset
 
     if local_dir.exists():
         return local_dir
 
-    response = requests.get(url)
-    if response.status_code == 200:
+    response = requests.get(url, headers=HEADERS)
+    if response.ok:
         files = response.json()
     else:
-        print(f"Failed to retrieve contents: {response.status_code}")
-        return None
+        raise ValueError(
+            f"Failed to retrieve contents ({response.status_code}): \n {response.text}"
+        )
 
     if not local_dir.exists():
         local_dir.mkdir(parents=True)
@@ -27,8 +39,16 @@ def download_map(dataset):
     for file in files:
         file_url = file["download_url"]
         file_name = file["name"]
-        file_response = requests.get(file_url)
+        file_response = requests.get(file_url, headers=HEADERS)
+        if not file_response.ok:
+            rmtree(local_dir)
+            raise ValueError(f"Failed to download file: {file_name}, \n{file_response.text}")
         with open(local_dir / file_name, "wb") as f:
             f.write(file_response.content)
 
     return local_dir
+
+
+if __name__ == "__main__":
+    download_map("naturalearth_lowres")
+    download_map("naturalearth_cities")