chore: fix tests & implement a semaphore on async downloads that was causing a throttle with too many files downloading at the same time

Author: luabida

.github/workflows/python-package.yml

@@ -71,10 +71,10 @@ jobs:
     - uses: actions/checkout@v4
 
     - name: Build Docker image
-      run: docker compose -f docker/docker-compose.yaml build
+      run: docker compose build
 
     - name: Start container
-      run: docker compose -f docker/docker-compose.yaml up -d
+      run: docker compose up -d
 
     - name: Wait for Jupyter
       run: |
@@ -84,8 +84,8 @@ jobs:
         done
 
     - name: Run tests inside container
-      run: docker compose -f docker/docker-compose.yaml exec -T -w /usr/src jupyter python3 -m pytest -vv pysus/tests/ --retries 3 --retry-delay 15 -x -o cache_dir=/tmp/.pytest_cache
+      run: docker compose exec -T -w /usr/src jupyter python3 -m pytest -vv pysus/tests/ --retries 3 --retry-delay 15 -x -o cache_dir=/tmp/.pytest_cache
 
     - name: Cleanup
       if: always()
-      run: docker compose -f docker/docker-compose.yaml down -v
+      run: docker compose down -v

.github/workflows/release.yaml

@@ -91,6 +91,11 @@ jobs:
       with:
         python-version: "3.12"
 
+    - name: Install system dependencies
+      run: |
+        sudo apt-get update
+        sudo apt-get install -y pandoc
+
     - name: Install dependencies
       run: |
         pip install poetry wget

Makefile

@@ -22,7 +22,7 @@ export PRINT_HELP_PYSCRIPT
 help:
 	@python -c "$$PRINT_HELP_PYSCRIPT" < $(MAKEFILE_LIST)
 
-DOCKER = docker compose -p pysus -f docker/docker-compose.yaml
+DOCKER = docker compose -p pysus
 SERVICE :=
 SEMANTIC_RELEASE = npx --yes \
           -p semantic-release \
@@ -56,7 +56,7 @@ test-pysus: ## run tests quickly with the default Python
 
 .PHONY: test-pysus-with-coverage
 test-pysus-with-coverage: ## run tests with coverage report
-	poetry run pytest -vv pysus/tests/ --retries 3 --retry-delay 15 --cov=pysus --cov-report=xml:coverage.xml --cov-report=term-missing
+	poetry run pytest -vv pysus/tests/ --cov=pysus --cov-report=xml:coverage.xml --cov-report=term-missing
 
 .PHONY: lint
 lint:

README.md

@@ -37,15 +37,15 @@ docker run -p 8888:8888 alertadengue/pysus
 Or build locally and start the container:
 
 ```bash
-docker compose -f docker/docker-compose.yaml up --build
+docker compose up --build
 ```
 
 Then open [http://127.0.0.1:8888/lab](http://127.0.0.1:8888/lab) in your browser.
 
 Stop the container:
 
 ```bash
-docker compose -f docker/docker-compose.yaml down
+docker compose down
 ```
 
 ## Quick Start
@@ -268,7 +268,7 @@ pytest tests/
 Run tests inside the Docker container:
 
 ```bash
-docker compose -f docker/docker-compose.yaml exec -T -w /usr/src jupyter python3 -m pytest pysus/tests/
+docker compose exec -T -w /usr/src jupyter python3 -m pytest pysus/tests/
 ```
 
 ## License

docs/source/conf.py

@@ -21,7 +21,10 @@
 }
 
 templates_path = ["_templates"]
+
+# Explicitly map extensions to ensure notebooks are routed to nbsphinx
 source_suffix = ".rst"
+
 master_doc = "index"
 
 project = "PySUS"

docs/source/databases/getting_started_pysus.ipynb

@@ -34,15 +34,15 @@ Or build locally and start the container:
 
 .. code-block:: bash
 
-   docker compose -f docker/docker-compose.yaml up --build
+   docker compose up --build
 
 Then open http://127.0.0.1:8888/lab in your browser.
 
 Stop the container with:
 
 .. code-block:: bash
 
-   docker compose -f docker/docker-compose.yaml down
+   docker compose down
 
 Development
 -----------

docs/source/installation.rst

@@ -8,12 +8,12 @@
 """
 
 import asyncio
-from typing import Literal
+from typing import Literal, cast
 
 import pandas as pd
 from pysus.api import types
 from pysus.api.client import PySUS
-from tqdm import tqdm
+from tqdm.asyncio import tqdm
 
 __all__ = [
     "sinan",
@@ -57,10 +57,9 @@ def _fetch_data(
     month : int | list[int], optional
         Month or list of months to fetch.
     show_progress : bool, optional
-        Whether to display a tqdm progress bar during download. Default is True.
+        Whether to display a tqdm progress bar during download.
     as_dataframe : bool, optional
         Whether to concatenate and return the data as a pandas DataFrame.
-        Default is False.
     **kwargs
         Additional arguments forwarded to :meth:`PySUS.read_parquet`.
 
@@ -71,48 +70,41 @@ def _fetch_data(
         as_dataframe is True, returns a concatenated DataFrame.
     """
 
-    async def _fetch():
-
+    async def _fetch() -> list[str] | pd.DataFrame:
         async with PySUS() as pysus:
-            years = [year] if isinstance(year, int) else (year or [None])
-            months = [month] if isinstance(month, int) else (month or [None])
+            files = await pysus.query(
+                dataset=dataset,
+                group=group,
+                state=state,
+                year=year,
+                month=month,
+            )
 
-            files = []
-            for y in years:
-                for m in months:
-                    files.extend(
-                        await pysus.query(
-                            dataset=dataset,
-                            group=group,
-                            state=state,
-                            year=y,
-                            month=m,
-                        )
-                    )
+            if not files:
+                return pd.DataFrame() if as_dataframe else cast(list[str], [])
+
+            sem = asyncio.Semaphore(3)
+
+            async def _throttled_download(f):
+                async with sem:
+                    return await pysus.download(f)
+
+            tasks = [_throttled_download(f) for f in files]
 
-            paths = []
             if show_progress:
-                for file in tqdm(
-                    files,
+                downloaded_files = await tqdm.gather(
+                    *tasks,
                     desc=f"Downloading {dataset}",
                     unit="file",
-                ):
-                    f = await pysus.download(file)
-                    paths.append(str(f.path))
+                )
             else:
-                for file in files:
-                    f = await pysus.download(file)
-                    paths.append(str(f.path))
+                downloaded_files = await asyncio.gather(*tasks)
+
+            paths: list[str] = [str(f.path) for f in downloaded_files]
 
             if as_dataframe:
-                return (
-                    pysus.read_parquet(
-                        paths,
-                        **kwargs,
-                    ).df()
-                    if paths
-                    else pd.DataFrame()
-                )
+                res = pysus.read_parquet(paths, **kwargs).df()
+                return cast(pd.DataFrame, res)
 
             return paths
 
@@ -132,9 +124,11 @@ async def _fetch():
                 "Install it with: pip install nest_asyncio"
             )
             raise RuntimeError(msg) from None
-        return loop.run_until_complete(_fetch())
-    else:
-        return asyncio.run(_fetch())
+        result = loop.run_until_complete(_fetch())
+        return cast(list[str] | pd.DataFrame, result)
+
+    result = asyncio.run(_fetch())
+    return cast(list[str] | pd.DataFrame, result)
 
 
 def sinan(

pysus/api/_impl/databases.py

@@ -314,7 +314,11 @@ async def download(
             )
             return await ExtensionFactory.instantiate(local_path)
 
-        except Exception as e:  # noqa: B902
+        except Exception as e:  # noqa
+            import traceback
+
+            traceback.print_exc()
+
             await self._update_state(
                 local_path,
                 str(remote_path),

pysus/api/client.py

@@ -239,7 +239,7 @@ def __init__(self, name: str, dataset_id: int, engine=None, **data) -> None:
         super().__init__(engine=engine, **data)
         self.dataset_name: str = name
         self.db_local: Path = self.cache_dir / f"catalog_{name}.duckdb"
-        self.db_remote: Path = Path(f"datasets/catalog_{name}.duckdb")
+        self.db_remote: Path = Path(f"public/catalog_{name}.duckdb")
         self.dataset_id = dataset_id