🎉 First commit !!!

backend/idfm_matrix_backend/models/line.py

@@ -0,0 +1,176 @@
+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