from asyncio import gather as asyncio_gather from collections import defaultdict from typing import Iterable, Self from sqlalchemy import ( BigInteger, Boolean, Column, Enum, ForeignKey, Integer, select, String, Table, ) from sqlalchemy.orm import Mapped, relationship, selectinload from sqlalchemy.orm.attributes import set_committed_value from sqlalchemy.sql.expression import tuple_ from ..db import Base, db from ..idfm_interface.idfm_types import ( IdfmState, IdfmLineState, TransportMode, TransportSubMode, ) from .stop import _Stop line_stop_association_table = Table( "line_stop_association_table", Base.metadata, Column("line_id", ForeignKey("lines.id")), Column("stop_id", ForeignKey("_stops.id")), ) class LinePicto(Base): db = db id = Column(String, primary_key=True) mime_type = Column(String, nullable=False) height_px = Column(Integer, nullable=False) width_px = Column(Integer, nullable=False) filename = Column(String, nullable=False) url = Column(String, nullable=False) thumbnail = Column(Boolean, nullable=False) format = Column(String, nullable=False) __tablename__ = "line_pictos" class Line(Base): db = db id = Column(String, primary_key=True) short_name = Column(String) name = Column(String, nullable=False) status = Column(Enum(IdfmLineState), nullable=False) transport_mode = Column(Enum(TransportMode), nullable=False) transport_submode = Column(Enum(TransportSubMode), nullable=False) network_name = Column(String) group_of_lines_id = Column(String) group_of_lines_shortname = Column(String) colour_web_hexa = Column(String, nullable=False) text_colour_hexa = Column(String, nullable=False) operator_id = Column(String) operator_name = Column(String) accessibility = Column(Enum(IdfmState), nullable=False) visual_signs_available = Column(Enum(IdfmState), nullable=False) audible_signs_available = Column(Enum(IdfmState), nullable=False) picto_id = Column(String, ForeignKey("line_pictos.id")) picto: Mapped[LinePicto] = relationship(LinePicto, lazy="selectin") record_id = Column(String, nullable=False) record_ts = Column(BigInteger, nullable=False) stops: Mapped[list["_Stop"]] = relationship( "_Stop", secondary=line_stop_association_table, back_populates="lines", lazy="selectin", ) __tablename__ = "lines" @classmethod async def get_by_name( cls, name: str, operator_name: None | str = None ) -> list[Self]: filters = {"name": name} if operator_name is not None: filters["operator_name"] = operator_name lines = None stmt = ( select(Line) .filter_by(**filters) .options(selectinload(Line.stops), selectinload(Line.picto)) ) res = await cls.db.session.execute(stmt) lines = res.scalars().all() return lines @classmethod async def _add_picto_to_line(cls, line: str | Self, picto: LinePicto) -> None: if isinstance(line, str): if (lines := await cls.get_by_name(line)) is not None: if len(lines) == 1: line = lines[0] else: for candidate_line in lines: if candidate_line.operator_name == "RATP": line = candidate_line break if isinstance(line, Line) and line.picto is None: line.picto = picto line.picto_id = picto.id @classmethod async def add_pictos(cls, line_to_pictos: dict[str | Self, LinePicto]) -> None: await asyncio_gather( *[ cls._add_picto_to_line(line, picto) for line, picto in line_to_pictos.items() ] ) await cls.db.session.commit() @classmethod async def add_stops(cls, line_to_stop_ids: Iterable[tuple[str, str, str]]) -> int: line_names_ops, stop_ids = set(), set() for line_name, operator_name, stop_id in line_to_stop_ids: line_names_ops.add((line_name, operator_name)) stop_ids.add(stop_id) res = await cls.db.session.execute( select(Line).where( tuple_(Line.name, Line.operator_name).in_(line_names_ops) ) ) lines = defaultdict(list) for line in res.scalars(): lines[(line.name, line.operator_name)].append(line) res = await cls.db.session.execute(select(_Stop).where(_Stop.id.in_(stop_ids))) stops = {stop.id: stop for stop in res.scalars()} found = 0 for line_name, operator_name, stop_id in line_to_stop_ids: if (stop := stops.get(stop_id)) is not None: if (stop_lines := lines.get((line_name, operator_name))) is not None: if len(stop_lines) > 1: print(stop_lines) for stop_line in stop_lines: stop_line.stops.append(stop) found += 1 else: print(f"No line found for {line_name}/{operator_name}") else: print( f"No stop found for {stop_id} id (used by {line_name}/{operator_name})" ) await cls.db.session.commit() return found