Skip to content

Commit

Permalink
Refactor progress bar to cli level
Browse files Browse the repository at this point in the history
  • Loading branch information
mfisher87 committed Jul 18, 2024
1 parent f8773f2 commit 06a98c8
Show file tree
Hide file tree
Showing 3 changed files with 17 additions and 19 deletions.
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ test = [
"pytest >=6",
"pytest-cov >=3",
"mypy >=1.10",
"types-tqdm",
]
dev = [
"pytest >=6",
Expand Down
9 changes: 5 additions & 4 deletions src/aross_stations_db/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

from aross_stations_db.config import CliLoadSettings
from aross_stations_db.db.setup import (
generate_event_objects,
generate_event_object,
load_events,
load_stations,
recreate_tables,
Expand Down Expand Up @@ -52,10 +52,11 @@ def init(skip_load: bool = False) -> None:
# The event processing steps are split into stages to provide better feadback at
# runtime. On slower systems, it can be unclear what the bottleneck is. In the
# long run, we should try to optimize this after learning more.
events = generate_event_objects(
tqdm(raw_events, desc="Reading events"),
)
events = [
generate_event_object(e) for e in tqdm(raw_events, desc="Reading events")
]

# TODO: Is there any way we can monitor this process with a progress bar?
logger.info("Loading events; this can take a minute or so")
load_events(events, session=db_session)
logger.info("Loaded events")
Expand Down
26 changes: 11 additions & 15 deletions src/aross_stations_db/db/setup.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import datetime as dt
from collections.abc import Iterator

from sqlalchemy import MetaData, insert
from sqlalchemy.orm import Session
Expand Down Expand Up @@ -65,20 +64,17 @@ def load_stations(stations: list[dict[str, str]], *, session: Session) -> None:
session.commit()


def generate_event_objects(raw_events: Iterator[dict[str, str]]) -> list[Event]:
return [
Event(
station_id=event["station_id"],
time_start=dt.datetime.fromisoformat(event["start"]),
time_end=dt.datetime.fromisoformat(event["end"]),
snow_on_ground=_snow_on_ground_status(event["sog"]),
rain_hours=int(event["RA"]),
freezing_rain_hours=int(event["FZRA"]),
solid_precipitation_hours=int(event["SOLID"]),
unknown_precipitation_hours=int(event["UP"]),
)
for event in raw_events
]
def generate_event_object(raw_event: dict[str, str]) -> Event:
return Event(
station_id=raw_event["station_id"],
time_start=dt.datetime.fromisoformat(raw_event["start"]),
time_end=dt.datetime.fromisoformat(raw_event["end"]),
snow_on_ground=_snow_on_ground_status(raw_event["sog"]),
rain_hours=int(raw_event["RA"]),
freezing_rain_hours=int(raw_event["FZRA"]),
solid_precipitation_hours=int(raw_event["SOLID"]),
unknown_precipitation_hours=int(raw_event["UP"]),
)


def load_events(events: list[Event], *, session: Session) -> None:
Expand Down

0 comments on commit 06a98c8

Please sign in to comment.