Move schedule_getter option to class constructor

lemonyte · Jun 9, 2024 · 5d82c18 · 5d82c18
1 parent 354e512
commit 5d82c18
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 11 deletions.
diff --git a/src/ferry_planner/route.py b/src/ferry_planner/route.py
@@ -250,15 +250,15 @@ def _find_routes_recurse(  # noqa: C901, PLR0912, PLR0913
 
 
 class RoutePlanBuilder:
-    def __init__(self, connection_db: ConnectionDB, /) -> None:
+    def __init__(self, *, connection_db: ConnectionDB, schedule_getter: ScheduleGetter) -> None:
         self._connection_db = connection_db
+        self._schedule_getter = schedule_getter
 
     async def make_route_plans(
         self,
         *,
         routes: Iterable[Route],
         options: RoutePlansOptions,
-        schedule_getter: ScheduleGetter,
     ) -> Sequence[RoutePlan]:
         route_plans = []
         for route in routes:
@@ -268,7 +268,6 @@ async def make_route_plans(
                 destination_index=1,
                 start_time=options.date.replace(hour=0, minute=0, second=0, microsecond=0),
                 options=options,
-                schedule_getter=schedule_getter,
             )
         return route_plans
 
@@ -280,7 +279,6 @@ async def _add_plan_segment(  # noqa: PLR0913
         destination_index: int,
         start_time: datetime,
         options: RoutePlansOptions,
-        schedule_getter: ScheduleGetter,
         segments: list[RoutePlanSegment] | None = None,
     ) -> bool:
         if segments is None:
@@ -304,7 +302,6 @@ async def _add_plan_segment(  # noqa: PLR0913
                     start_time=start_time,
                     options=options,
                     connection=connection,
-                    schedule_getter=schedule_getter,
                 )
             if isinstance(connection, CarConnection):
                 driving_duration_limit = 6 * 60 * 60
@@ -330,7 +327,6 @@ async def _add_plan_segment(  # noqa: PLR0913
                     segments=segments,
                     start_time=arrive_time,
                     options=options,
-                    schedule_getter=schedule_getter,
                 )
         finally:
             delete_start = destination_index - 1
@@ -347,12 +343,11 @@ async def _add_ferry_connection(  # noqa: C901, PLR0912, PLR0913
         start_time: datetime,
         options: RoutePlansOptions,
         connection: FerryConnection,
-        schedule_getter: ScheduleGetter,
     ) -> bool:
         result = False
         depature_terminal = connection.origin
         day = start_time.replace(hour=0, minute=0, second=0, microsecond=0)
-        schedule = await schedule_getter(connection.origin.id, connection.destination.id, date=day)
+        schedule = await self._schedule_getter(connection.origin.id, connection.destination.id, date=day)
         if not schedule:
             return False
         for sailing in schedule.sailings:
@@ -420,7 +415,6 @@ async def _add_ferry_connection(  # noqa: C901, PLR0912, PLR0913
                 segments=segments,
                 start_time=arrive_time,
                 options=options,
-                schedule_getter=schedule_getter,
             )
             if recursion_result is False:
                 break

diff --git a/src/ferry_planner/server.py b/src/ferry_planner/server.py
@@ -53,7 +53,10 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
     refresh_interval=config.schedules.refresh_interval,
 )
 route_builder = RouteBuilder(connection_db)
-route_plan_builder = RoutePlanBuilder(connection_db)
+route_plan_builder = RoutePlanBuilder(
+    connection_db=connection_db,
+    schedule_getter=schedule_db.get,
+)
 app = FastAPI(lifespan=lifespan)
 app.mount("/static", StaticFiles(directory=ROOT_DIR / "static"), name="static")
 templates = Jinja2Templates(directory=ROOT_DIR / "templates")
@@ -100,7 +103,6 @@ async def api_routeplans(options: RoutePlansOptions) -> Sequence[RoutePlan]:
         await route_plan_builder.make_route_plans(
             routes=routes,
             options=options,
-            schedule_getter=schedule_db.get,
         ),
     )
     route_plans.sort(key=lambda plan: plan.duration)