Configurable progressbars

src/traffic/__init__.py

@@ -12,7 +12,7 @@
 from . import drawing  # noqa: F401
 
 __version__ = version("traffic")
-__all__ = ["config_dir", "config_file", "cache_dir"]
+__all__ = ["config_dir", "config_file", "cache_dir", "tqdm_style"]
 
 # Set up the library root logger
 _log = logging.getLogger(__name__)
@@ -65,6 +65,10 @@
 if not cache_dir.exists():
     cache_dir.mkdir(parents=True)
 
+# -- Tqdm Style Configuration --
+tqdm_style = config.get("global", "tqdm_style", fallback="auto")
+_log.info(f"Selected tqdm style: {tqdm_style}")
+
 # -- Plugin management --
 
 _enabled_plugins_raw = config.get("plugins", "enabled_plugins", fallback="")

src/traffic/algorithms/cpa.py

@@ -11,12 +11,10 @@
     cast,
 )
 
-from tqdm.rich import tqdm
-
 import pandas as pd
 import pyproj
 
-from ..core import Flight
+from ..core import Flight, tqdm
 from ..core.mixins import DataFrameMixin
 
 if TYPE_CHECKING:

src/traffic/core/__init__.py

@@ -1,4 +1,4 @@
-# ruff: noqa: I001
+# ruff: noqa: I001, E402
 """
 It is crucial that the imports do not change order,
 hence the following line:
@@ -8,7 +8,35 @@
 import logging
 import sys
 from types import TracebackType
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Iterable, Iterator, Optional, TypeVar
+
+from tqdm.auto import tqdm as _tqdm_auto
+from tqdm.autonotebook import tqdm as _tqdm_autonotebook
+from tqdm.rich import tqdm as _tqdm_rich
+from tqdm.std import tqdm as _tqdm_std
+
+from .. import tqdm_style
+from .types import ProgressbarType
+
+T = TypeVar("T")
+
+
+def silent_tqdm(
+    iterable: Iterable[T], *args: Any, **kwargs: Any
+) -> Iterator[T]:
+    yield from iterable  # Dummy tqdm function
+
+
+tqdm_dict: Dict[str, ProgressbarType] = {
+    "autonotebook": _tqdm_autonotebook,
+    "auto": _tqdm_auto,
+    "rich": _tqdm_rich,
+    "silent": silent_tqdm,
+    "std": _tqdm_std,
+}
+
+
+tqdm = tqdm_dict[tqdm_style]
 
 # WARNING!! Don't change order of import in this file
 from .flight import Flight
@@ -30,6 +58,7 @@
     "LazyTraffic",
     "loglevel",
     "faulty_flight",
+    "tqdm",
 ]
 
 

src/traffic/core/flight.py

@@ -2676,7 +2676,7 @@ def query_ehs(
         self,
         data: Union[None, pd.DataFrame, "RawData"] = None,
         failure_mode: str = "info",
-        progressbar: Union[bool, tt.ProgressbarType[Any]] = False,
+        progressbar: Union[bool, tt.ProgressbarType] = False,
     ) -> Flight:
         """Extends data with extra columns from EHS messages.
 

src/traffic/core/lazy.py

@@ -24,11 +24,11 @@
 )
 
 import astor
-from tqdm.rich import tqdm
 
 import numpy as np
 import pandas as pd
 
+from . import tqdm
 from .flight import Flight
 from .mixins import GeographyMixin
 

src/traffic/core/traffic.py

@@ -29,7 +29,6 @@
 from ipyleaflet import Map as LeafletMap
 from ipywidgets import HTML
 from rich.console import Console, ConsoleOptions, RenderResult
-from tqdm.rich import tqdm
 
 import numpy as np
 import pandas as pd
@@ -42,6 +41,7 @@
 from ..core.cache import property_cache
 from ..core.structure import Airport
 from ..core.time import time_or_delta, timelike, to_datetime
+from . import tqdm
 from .flight import Flight
 from .intervals import Interval, IntervalCollection
 from .lazy import LazyTraffic, lazy_evaluation

src/traffic/core/types.py

@@ -1,14 +1,20 @@
 from __future__ import annotations
 
-from typing import Callable, Iterable, TypeVar
+from typing import Any, Iterable, Iterator, TypeVar
 
 from typing_extensions import Annotated, Protocol
 
 import numpy as np
 import numpy.typing as npt
 
 T = TypeVar("T")
-ProgressbarType = Callable[[Iterable[T]], Iterable[T]]
+
+
+class ProgressbarType(Protocol):
+    def __call__(
+        self, iterable: Iterable[T], *args: Any, **kwargs: Any
+    ) -> Iterator[T]:
+        ...
 
 
 class HasBounds(Protocol):

src/traffic/data/adsb/decode.py

@@ -27,19 +27,18 @@
 )
 
 import pyModeS as pms
-from tqdm.rich import tqdm
 
 import numpy as np
 import pandas as pd
 
-from ...core import Flight, Traffic
+from ...core import Flight, Traffic, tqdm
 from ...core.mixins import DataFrameMixin
 from ...core.types import ProgressbarType
 from ...data.basic.airports import Airport
 
 D = TypeVar("D", bound="ModeS_Decoder")
 R = TypeVar("R", bound="RawData")
-U = Tuple[int, Tuple[Any, Any]]
+T = TypeVar("T")
 
 _log = logging.getLogger(__name__)
 
@@ -1782,7 +1781,7 @@ def decode(
         uncertainty: bool = False,
         # assume that data you get from a file or DB is already CRC-checked
         crc_check: bool = False,
-        progressbar: Union[bool, ProgressbarType[U]] = True,
+        progressbar: Union[bool, ProgressbarType] = True,
         progressbar_kw: Optional[Dict[str, Any]] = None,
         redefine_mag: int = 10,
     ) -> Optional[Traffic]:
@@ -1793,16 +1792,20 @@ def decode(
             if progressbar_kw is None:
                 progressbar_kw = dict()
 
-            def custom_tqdm(x: Iterable[U]) -> Iterable[U]:
-                return tqdm(  # type: ignore
-                    x, total=self.data.shape[0], **progressbar_kw
+            def custom_tqdm(
+                iterable: Iterable[T], *args: Any, **kwargs: Any
+            ) -> Iterator[T]:
+                yield from tqdm(
+                    iterable, total=self.data.shape[0], **progressbar_kw
                 )
 
             progressbar = custom_tqdm
         elif progressbar is False:
 
-            def identity(x: Iterable[U]) -> Iterable[U]:
-                return x
+            def identity(
+                iterable: Iterable[T], *args: Any, **kwargs: Any
+            ) -> Iterator[T]:
+                yield from iterable
 
             progressbar = identity
 

src/traffic/data/adsb/flarm.py

@@ -2,11 +2,11 @@
 
 from pitot.geodesy import distance
 from pyModeS.decoder.flarm import DecodedMessage, flarm
-from tqdm.rich import tqdm
 from typing_extensions import Annotated
 
 import pandas as pd
 
+from ...core import tqdm
 from ...core.mixins import DataFrameMixin
 from ...core.traffic import Traffic
 

src/traffic/data/adsb/opensky.py

@@ -517,7 +517,7 @@ def flarm(
         compress: bool = False,
         limit: None | int = None,
         other_params: str = "",
-        progressbar: bool | ProgressbarType[Any] = True,
+        progressbar: bool | ProgressbarType = True,
     ) -> None | FlarmData:
         df = self.impala_client.flarm(
             start,

src/traffic/data/basic/aircraft.py

@@ -11,10 +11,10 @@
 from typing import Any, ClassVar, Dict, TypeVar
 
 import rich.repr
-from tqdm.rich import tqdm
 
 import pandas as pd
 
+from ...core import tqdm
 from ...core.mixins import DataFrameMixin, FormatMixin
 
 _log = logging.getLogger(__name__)

src/traffic/data/basic/airports.py

@@ -6,11 +6,11 @@
 from typing import Any, ClassVar
 
 import requests
-from tqdm.rich import tqdm
 
 import pandas as pd
 
 from ... import cache_expiration
+from ...core import tqdm
 from ...core.mixins import GeoDBMixin
 from ...core.structure import Airport
 

src/traffic/data/basic/runways.py

@@ -17,14 +17,14 @@
 
 import requests
 from pitot.geodesy import bearing, destination
-from tqdm.rich import tqdm
 
 import numpy as np
 import pandas as pd
 from shapely.geometry import base, shape
 from shapely.ops import linemerge
 
 from ... import cache_expiration
+from ...core import tqdm
 from ...core.mixins import DataFrameMixin, HBoxMixin, PointMixin, ShapelyMixin
 
 if TYPE_CHECKING:

src/traffic/data/datasets/__init__.py

@@ -2,12 +2,11 @@
 from hashlib import md5
 from typing import Any, Dict
 
-from tqdm.rich import tqdm
-
 import pandas as pd
 
 from ... import cache_dir, config
 from ...core import Traffic
+from ...core.progress_bar import tqdm
 
 _squawk7700_url = "https://zenodo.org/record/3937483/"
 

src/traffic/data/eurocontrol/aixm/airspaces.py

@@ -8,12 +8,12 @@
 import geopandas as gpd
 from geopandas.geodataframe import GeoDataFrame
 from lxml import etree
-from tqdm.rich import tqdm
 
 import pandas as pd
 from shapely.geometry import Polygon
 from shapely.ops import orient, unary_union
 
+from ....core import tqdm
 from ....core.airspace import (
     Airspace,
     Airspaces,

src/traffic/data/samples/__init__.py

@@ -6,9 +6,8 @@
 from typing import Any, Union, cast
 
 import pandas as pd
-from traffic.core import Airspace
 
-from ...core import Flight, Traffic
+from ...core import Airspace, Flight, Traffic
 from ..eurocontrol.ddr.so6 import SO6
 
 _current_dir = Path(__file__).parent

src/traffic/traffic.conf

@@ -14,6 +14,13 @@
 aixm_path =
 nm_path =
 
+## Set the style of progress bars
+## Progress bar are based on tqdm, which comes in different flavours.
+## - 'silent' disable progressbars
+## - 'default' option is 'auto'
+
+tqdm_style = auto
+
 [opensky]
 
 ## Input here your credentials for accessing the OpenSky Impala shell
@@ -63,6 +70,7 @@ expiration = 90 days
 ## Files are removed when the library is imported.
 purge = 180 days
 
+
 [plugins]
 
 ## input here a list of plugins you wish to activate each time you import