feat: add automatic download progress bar removal (#183)

* chore: add automatic terminal detection * chore: add changelog entry * fix(pre-commit.ci): auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * chore: add automatic terminal detection * chore: ignore refurb suggestions * ci: add terminal mode to tests * test: add debug print * ci: add env variables passing to tox * test: remove debug print * ci: change conftest fixture order --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
kraina-ai · Jan 3, 2025 · 68d8d97 · 68d8d97
1 parent fa3c03a
commit 68d8d97
Show file tree

Hide file tree

Showing 11 changed files with 74 additions and 51 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Automatic download progress bar hiding when verbosity is set to `silent`.
+
 ## [0.12.0] - 2024-11-03
 
 ### Added

diff --git a/quackosm/_constants.py b/quackosm/_constants.py
@@ -1,7 +1,11 @@
 """Constants used across the project."""
 
+import os
+
 WGS84_CRS = "EPSG:4326"
 
 FEATURES_INDEX = "feature_id"
 
 GEOMETRY_COLUMN = "geometry"
+
+FORCE_TERMINAL = os.getenv("FORCE_TERMINAL_MODE", "false").lower() == "true"
diff --git a/quackosm/_rich_progress.py b/quackosm/_rich_progress.py
@@ -2,7 +2,6 @@
 """Wrapper over Rich progress bar."""
 
 import json
-import os
 import time
 from collections.abc import Iterable
 from datetime import timedelta
@@ -25,6 +24,8 @@
     TimeRemainingColumn,
 )
 
+from quackosm._constants import FORCE_TERMINAL
+
 __all__ = ["TaskProgressSpinner", "TaskProgressBar"]
 
 TOTAL_STEPS = 32
@@ -234,12 +235,10 @@ def __init__(
             self.major_steps_prefix = ""
 
         if not self.verbosity_mode == "silent":
-            self.force_terminal = os.getenv("FORCE_TERMINAL_MODE", "false").lower() == "true"
-
             self.console = Console(
-                force_interactive=False if self.force_terminal else None,
-                force_jupyter=False if self.force_terminal else None,
-                force_terminal=True if self.force_terminal else None,
+                force_interactive=False if FORCE_TERMINAL else None,  # noqa: FURB110
+                force_jupyter=False if FORCE_TERMINAL else None,  # noqa: FURB110
+                force_terminal=True if FORCE_TERMINAL else None,  # noqa: FURB110
             )
             self.transient_progress_cls = TransientProgress
 

diff --git a/quackosm/conftest.py b/quackosm/conftest.py
@@ -2,7 +2,6 @@
 
 import doctest
 import shutil
-import urllib.request
 from doctest import OutputChecker
 from pathlib import Path
 
@@ -12,6 +11,7 @@
 from pooch import get_logger as get_pooch_logger
 from pooch import retrieve
 
+from quackosm._constants import FORCE_TERMINAL
 from quackosm.osm_extracts.extract import OsmExtractSource
 from quackosm.osm_extracts.geofabrik import _get_geofabrik_index
 
@@ -35,25 +35,6 @@ def check_output(self: doctest.OutputChecker, want: str, got: str, optionflags:
 EXTRACTS_NAMES = ["monaco", "kiribati", "maldives"]
 
 
-@pytest.fixture(autouse=True, scope="session")
-def add_pbf_files(doctest_namespace):  # type: ignore
-    """Download PBF files used in doctests."""
-    download_directory = Path("files")
-    download_directory.mkdir(parents=True, exist_ok=True)
-
-    geofabrik_index = _get_geofabrik_index()
-    for extract_name in EXTRACTS_NAMES:
-        pbf_file_download_url = LFS_DIRECTORY_URL + f"{extract_name}-latest.osm.pbf"
-        pbf_file_path = download_directory / f"{extract_name}.osm.pbf"
-        geofabrik_download_path = geofabrik_index[geofabrik_index["name"] == extract_name].iloc[0][
-            "file_name"
-        ]
-        geofabrik_pbf_file_path = download_directory / f"{geofabrik_download_path}.osm.pbf"
-        urllib.request.urlretrieve(pbf_file_download_url, pbf_file_path)
-        doctest_namespace[f"{extract_name}_pbf_path"] = pbf_file_path
-        shutil.copy(pbf_file_path, geofabrik_pbf_file_path)
-
-
 @pytest.fixture(autouse=True, scope="session")
 def download_osm_extracts_indexes():  # type: ignore
     """Download OSM extract indexes files to cache."""
@@ -74,9 +55,34 @@ def download_osm_extracts_indexes():  # type: ignore
             file_download_url,
             fname=file_name,
             path=download_directory,
-            progressbar=True,
+            progressbar=not FORCE_TERMINAL,
+            known_hash=None,
+        )
+
+
+@pytest.fixture(autouse=True, scope="session")
+def add_pbf_files(doctest_namespace, download_osm_extracts_indexes):  # type: ignore
+    """Download PBF files used in doctests."""
+    download_directory = Path("files")
+    download_directory.mkdir(parents=True, exist_ok=True)
+
+    geofabrik_index = _get_geofabrik_index()
+    for extract_name in EXTRACTS_NAMES:
+        pbf_file_download_url = LFS_DIRECTORY_URL + f"{extract_name}-latest.osm.pbf"
+        pbf_file_path = download_directory / f"{extract_name}.osm.pbf"
+        geofabrik_download_path = geofabrik_index[geofabrik_index["name"] == extract_name].iloc[0][
+            "file_name"
+        ]
+        geofabrik_pbf_file_path = download_directory / f"{geofabrik_download_path}.osm.pbf"
+        retrieve(
+            pbf_file_download_url,
+            fname=f"{extract_name}.osm.pbf",
+            path=download_directory,
+            progressbar=not FORCE_TERMINAL,
             known_hash=None,
         )
+        doctest_namespace[f"{extract_name}_pbf_path"] = pbf_file_path
+        shutil.copy(pbf_file_path, geofabrik_pbf_file_path)
 
 
 @pytest.fixture(autouse=True, scope="session")

diff --git a/quackosm/functions.py b/quackosm/functions.py
@@ -591,7 +591,7 @@ def convert_osm_extract_to_duckdb(
         'files/geofabrik_europe_monaco_nofilter_noclip_compact.duckdb'
     """
     downloaded_osm_extract = download_extract_by_query(
-        query=osm_extract_query, source=osm_extract_source
+        query=osm_extract_query, source=osm_extract_source, progressbar=verbosity_mode != "silent"
     )
     return PbfFileReader(
         tags_filter=tags_filter,
@@ -1176,7 +1176,7 @@ def convert_osm_extract_to_parquet(
         'files/geofabrik_europe_monaco_nofilter_noclip_compact.parquet'
     """
     downloaded_osm_extract = download_extract_by_query(
-        query=osm_extract_query, source=osm_extract_source
+        query=osm_extract_query, source=osm_extract_source, progressbar=verbosity_mode != "silent"
     )
     return PbfFileReader(
         tags_filter=tags_filter,
@@ -1668,7 +1668,7 @@ def convert_osm_extract_to_geodataframe(
         [7906 rows x 2 columns]
     """
     downloaded_osm_extract = download_extract_by_query(
-        query=osm_extract_query, source=osm_extract_source
+        query=osm_extract_query, source=osm_extract_source, progressbar=verbosity_mode != "silent"
     )
     return PbfFileReader(
         tags_filter=tags_filter,

diff --git a/quackosm/osm_extracts/__init__.py b/quackosm/osm_extracts/__init__.py
@@ -6,7 +6,6 @@
 """
 
 import difflib
-import os
 import warnings
 from collections.abc import Iterable
 from functools import partial
@@ -24,6 +23,7 @@
 from shapely.geometry.base import BaseGeometry, BaseMultipartGeometry
 from tqdm.contrib.concurrent import process_map
 
+from quackosm._constants import FORCE_TERMINAL
 from quackosm._exceptions import (
     GeometryNotCoveredError,
     GeometryNotCoveredWarning,
@@ -54,14 +54,15 @@
 
 
 def download_extracts_pbf_files(
-    extracts: list[OpenStreetMapExtract], download_directory: Path
+    extracts: list[OpenStreetMapExtract], download_directory: Path, progressbar: bool = True
 ) -> list[Path]:
     """
     Download OSM extracts as PBF files.
 
     Args:
         extracts (list[OpenStreetMapExtract]): List of extracts to download.
         download_directory (Path): Directory where PBF files should be saved.
+        progressbar (bool, optional): Show progress bar. Defaults to True.
 
     Returns:
         list[Path]: List of downloaded file paths.
@@ -74,7 +75,7 @@ def download_extracts_pbf_files(
             extract.url,
             fname=f"{extract.file_name}.osm.pbf",
             path=download_directory,
-            progressbar=True,
+            progressbar=progressbar and not FORCE_TERMINAL,
             known_hash=None,
         )
         downloaded_extracts_paths.append(Path(file_path))
@@ -194,17 +195,26 @@ def get_extract_by_query(
 
 
 @overload
-def download_extract_by_query(query: str) -> Path: ...
+def download_extract_by_query(
+    query: str, *, download_directory: Union[str, Path] = "files", progressbar: bool = True
+) -> Path: ...
 
 
 @overload
-def download_extract_by_query(query: str, source: Union[OsmExtractSource, str]) -> Path: ...
+def download_extract_by_query(
+    query: str,
+    source: Union[OsmExtractSource, str],
+    *,
+    download_directory: Union[str, Path] = "files",
+    progressbar: bool = True,
+) -> Path: ...
 
 
 def download_extract_by_query(
     query: str,
     source: Union[OsmExtractSource, str] = "any",
     download_directory: Union[str, Path] = "files",
+    progressbar: bool = True,
 ) -> Path:
     """
     Download an OSM extract by name.
@@ -215,12 +225,13 @@ def download_extract_by_query(
             'BBBike', 'OSM_fr'. Defaults to 'any'.
         download_directory (Union[str, Path], optional): Directory where the file should be
             downloaded. Defaults to "files".
+        progressbar (bool, optional): Show progress bar. Defaults to True.
 
     Returns:
         Path: Path to the downloaded OSM extract.
     """
     matching_extract = get_extract_by_query(query, source)
-    return download_extracts_pbf_files([matching_extract], Path(download_directory))[0]
+    return download_extracts_pbf_files([matching_extract], Path(download_directory), progressbar)[0]
 
 
 def display_available_extracts(
@@ -507,14 +518,13 @@ def _find_smallest_containing_extracts(
             allow_uncovered_geometry=allow_uncovered_geometry,
         )
 
-        force_terminal = os.getenv("FORCE_TERMINAL_MODE", "false").lower() == "true"
         for extract_ids_list in process_map(
             find_extracts_func,
             geometries,
             desc="Finding matching extracts",
             max_workers=num_of_multiprocessing_workers,
             chunksize=ceil(total_polygons / (4 * num_of_multiprocessing_workers)),
-            disable=True if force_terminal else False,
+            disable=FORCE_TERMINAL,
         ):
             unique_extracts_ids.update(extract_ids_list)
     else:
@@ -727,14 +737,13 @@ def _filter_extracts(
             sorted_extracts_gdf=sorted_extracts_gdf,
         )
 
-        force_terminal = os.getenv("FORCE_TERMINAL_MODE", "false").lower() == "true"
         for extract_ids_list in process_map(
             filter_extracts_func,
             geometries,
             desc="Filtering extracts",
             max_workers=num_of_multiprocessing_workers,
             chunksize=ceil(total_geometries / (4 * num_of_multiprocessing_workers)),
-            disable=True if force_terminal else False,
+            disable=FORCE_TERMINAL,
         ):
             filtered_extracts_ids.update(extract_ids_list)
     else:

diff --git a/quackosm/osm_extracts/bbbike.py b/quackosm/osm_extracts/bbbike.py
@@ -4,13 +4,13 @@
 This module contains wrapper for publically available BBBike download server.
 """
 
-import os
 from typing import Optional
 
 import geopandas as gpd
 import requests
 from tqdm import tqdm
 
+from quackosm._constants import FORCE_TERMINAL
 from quackosm.osm_extracts._poly_parser import parse_polygon_file
 from quackosm.osm_extracts.extract import (
     OpenStreetMapExtract,
@@ -71,12 +71,10 @@ def _iterate_bbbike_index() -> list[OpenStreetMapExtract]:  # pragma: no cover
         if extract_href.text != ".."
     ]
 
-    force_terminal = os.getenv("FORCE_TERMINAL_MODE", "false").lower() == "true"
-
     bbbike_enum_value = OsmExtractSource.bbbike.value
 
     with tqdm(
-        disable=True if force_terminal else False, desc=bbbike_enum_value, total=len(extract_names)
+        disable=FORCE_TERMINAL, desc=bbbike_enum_value, total=len(extract_names)
     ) as pbar:
         for extract_name in extract_names:
             pbar.set_description(f"{bbbike_enum_value}_{extract_name}")

diff --git a/quackosm/osm_extracts/osm_fr.py b/quackosm/osm_extracts/osm_fr.py
@@ -4,14 +4,14 @@
 This module contains wrapper for publically available OpenStreetMap.fr download server.
 """
 
-import os
 import re
 from typing import Any, Optional
 
 import geopandas as gpd
 import requests
 from tqdm import tqdm
 
+from quackosm._constants import FORCE_TERMINAL
 from quackosm.osm_extracts._poly_parser import parse_polygon_file
 from quackosm.osm_extracts.extract import (
     OpenStreetMapExtract,
@@ -44,9 +44,8 @@ def _load_openstreetmap_fr_index() -> gpd.GeoDataFrame:  # pragma: no cover
     Returns:
         gpd.GeoDataFrame: Extracts index with metadata.
     """
-    force_terminal = os.getenv("FORCE_TERMINAL_MODE", "false").lower() == "true"
     extracts = []
-    with tqdm(disable=True if force_terminal else False) as pbar:
+    with tqdm(disable=FORCE_TERMINAL) as pbar:
         extract_soup_objects = _gather_all_openstreetmap_fr_urls(
             OsmExtractSource.osm_fr.value, "/", pbar
         )

diff --git a/quackosm/pbf_file_reader.py b/quackosm/pbf_file_reader.py
@@ -38,7 +38,7 @@
 from shapely.geometry import LinearRing, Polygon
 from shapely.geometry.base import BaseGeometry, BaseMultipartGeometry
 
-from quackosm._constants import FEATURES_INDEX, GEOMETRY_COLUMN, WGS84_CRS
+from quackosm._constants import FEATURES_INDEX, FORCE_TERMINAL, GEOMETRY_COLUMN, WGS84_CRS
 from quackosm._exceptions import (
     EmptyResultWarning,
     InvalidGeometryFilter,
@@ -593,7 +593,9 @@ def convert_geometry_to_parquet(
             geometry_coverage_iou_threshold=self.geometry_coverage_iou_threshold,
             allow_uncovered_geometry=self.allow_uncovered_geometry,
         )
-        pbf_files = download_extracts_pbf_files(matching_extracts, self.working_directory)
+        pbf_files = download_extracts_pbf_files(
+            matching_extracts, self.working_directory, progressbar=self.verbosity_mode != "silent"
+        )
         return self.convert_pbf_to_parquet(
             pbf_files,
             result_file_path=result_file_path,
@@ -965,7 +967,7 @@ def _parse_pbf_file(
                 pbf_path,
                 fname=Path(pbf_path).name,
                 path=self.working_directory,
-                progressbar=True,
+                progressbar=self.verbosity_mode != "silent" and not FORCE_TERMINAL,
                 known_hash=None,
             )
 

diff --git a/tests/benchmark/test_big_file.py b/tests/benchmark/test_big_file.py
@@ -8,6 +8,7 @@
 from pooch import retrieve
 
 from quackosm import PbfFileReader, geocode_to_geometry
+from quackosm._constants import FORCE_TERMINAL
 from quackosm._osm_tags_filters import OsmTagsFilter
 
 
@@ -29,7 +30,7 @@ def test_big_file(extract_name: str, geocode_filter: list[str], tags_filter: Osm
         f"https://download.geofabrik.de/europe/{extract_name}-latest.osm.pbf",
         fname=f"{extract_name}.osm.pbf",
         path=files_dir,
-        progressbar=True,
+        progressbar=not FORCE_TERMINAL,
         known_hash=None,
     )
 

diff --git a/tox.ini b/tox.ini
@@ -14,6 +14,7 @@ groups =
 deps =
     coverage
     pre-commit
+passenv = *
 commands =
     coverage run --data-file=.coverage.doc.tests --source=quackosm -m pytest -v -s --durations=20 --doctest-modules --doctest-continue-on-failure quackosm
     coverage run --data-file=.coverage.base.tests --source=quackosm -m pytest -v -s --durations=20 tests/base