metro/db_graph_loader.py - master - Flash_A/SubjectTemplate

metro/db_graph_loader.py
194 строки8 KB
FlashAnton
Добавлены материалы от студентов
22 янв 2026, 21:16
22 янв 2026, 21:16511daf0
"""
Загрузчик данных графика движения поездов из базы данных PostgreSQL
Заменитель excel_graph_loader.py
"""
from typing import Dict, List, Any, Optional
from datetime import time

from flask import current_app

# Импортируем модели из app для избежания циклических зависимостей
def get_models():
    """Получает модели и объект db из текущего приложения"""
    from app import Station, Direction, HourInterval, DirectionHourCount, TrainRun
    db = current_app.extensions['sqlalchemy'].db
    return db, Station, Direction, HourInterval, DirectionHourCount, TrainRun

BASE_HOUR = 5  # 05:00 local start for график


class DatabaseGraphError(RuntimeError):
    """Raised when database data cannot be loaded."""


def _format_time(offset_minutes: float) -> str:
    """Форматирует время из минут от базового часа в формат HH:MM"""
    total_minutes = BASE_HOUR * 60 + offset_minutes
    total_minutes = total_minutes % (24 * 60)
    hours = int(total_minutes // 60)
    minutes = int(round(total_minutes % 60))
    return f"{hours:02d}:{minutes:02d}"


def _generate_runs_from_db(
    direction, hour_counts: List[Dict[str, Any]], travel_minutes: float
) -> List[Dict[str, Any]]:
    """Генерирует список рейсов на основе данных из БД"""
    runs: List[Dict[str, Any]] = []
    counter = 1
    
    for hour_data in hour_counts:
        hour_interval = hour_data['hour_interval']
        trips = hour_data['train_count']
        if not trips:
            continue
        
        # Вычисляем временные границы интервала
        start_hour = hour_interval.start_hour
        end_hour = hour_interval.end_hour
        
        # Нормализуем часы относительно базового часа
        while start_hour < BASE_HOUR:
            start_hour += 24
        while end_hour <= start_hour:
            end_hour += 24
        
        start_min = (start_hour - BASE_HOUR) * 60
        end_min = (end_hour - BASE_HOUR) * 60
        
        # Распределяем поезда равномерно по интервалу
        interval = 60 / trips
        for trip_idx in range(trips):
            departure_min = start_min + (trip_idx + 0.5) * interval
            arrival_min = departure_min + travel_minutes
            
            prefix = f"{direction.code}_{direction.start_idx}_{direction.end_idx}"
            train_id = f"{prefix}-{counter:03d}"
            
            runs.append({
                "train_id": train_id,
                "departure_minutes": round(departure_min, 2),
                "arrival_minutes": round(arrival_min, 2),
                "departure_time": _format_time(departure_min),
                "arrival_time": _format_time(arrival_min),
            })
            counter += 1
    
    return runs


def load_database_graph_data() -> Dict[str, Any]:
    STATION_DISTANCES = {
        'Москва': 0,
        'Тверь': 167,
        'Бологое': 319,
        'Санкт-Петербург': 650
    }

    """
    Загружает данные графика движения поездов из базы данных PostgreSQL.
    Возвращает данные в том же формате, что и load_excel_graph_data.
    """
    try:
        db, Station, Direction, HourInterval, DirectionHourCount, TrainRun = get_models()

        # Загружаем станции
        stations = db.session.scalars(db.select(Station).order_by(Station.order_index)).all()
        if not stations:
            raise DatabaseGraphError("В базе данных нет станций. Загрузите данные станций.")

        station_names = [s.name for s in stations]
        station_distances = [STATION_DISTANCES.get(s.name, 0) for s in stations]

        # Загружаем часовые интервалы
        hour_intervals = db.session.scalars(db.select(HourInterval).order_by(HourInterval.order_index)).all()
        if not hour_intervals:
            raise DatabaseGraphError("В базе данных нет часовых интервалов. Загрузите данные интервалов.")

        hour_labels = [hi.label for hi in hour_intervals]

        # Загружаем направления
        directions_data = []
        directions = db.session.scalars(db.select(Direction)).all()

        for direction in directions:
            hour_counts_data = db.session.scalars(db.select(DirectionHourCount).filter_by(direction_id=direction.id).join(HourInterval).order_by(HourInterval.order_index)).all()

            if not hour_counts_data or sum(hc.train_count for hc in hour_counts_data) == 0:
                continue

            hour_counts = [{'hour_interval': hc.hour_interval, 'train_count': hc.train_count} for hc in hour_counts_data]
            travel_minutes = direction.travel_minutes or 51.0

            # Загрузка или генерация рейсов
            runs = []
            existing_runs = db.session.scalars(db.select(TrainRun).filter_by(direction_id=direction.id).order_by(TrainRun.departure_minutes)).all()
            if existing_runs:
                for run in existing_runs:
                    runs.append({
                        "train_id": run.train_id,
                        "departure_minutes": run.departure_minutes,
                        "arrival_minutes": run.arrival_minutes,
                        "departure_time": run.departure_time.strftime('%H:%M'),
                        "arrival_time": run.arrival_time.strftime('%H:%M'),
                    })
            else:
                runs = _generate_runs_from_db(direction, hour_counts, travel_minutes)

            train_type = "Обычный"
            color = "blue"
            if travel_minutes <= 240:
                train_type = "Сапсан"
                color = "red"
            elif travel_minutes > 480:
                train_type = "Грузовой"
                color = "green"

            directions_data.append({
                "name": direction.name,
                "start_station": direction.start_station.name,
                "end_station": direction.end_station.name,
                "start_idx": direction.start_idx,
                "end_idx": direction.end_idx,
                "color": color,
                "code": direction.code,
                "travel_minutes": travel_minutes,
                "hour_counts": [hc['train_count'] for hc in hour_counts],
                "total_departures": sum(hc['train_count'] for hc in hour_counts),
                "runs": runs,
                "train_type": train_type
            })

        # Вычисляем метаданные временной линии
        if hour_intervals:
            first_interval = hour_intervals[0]
            last_interval = hour_intervals[-1]
            start_hour = first_interval.start_hour
            while start_hour < BASE_HOUR:
                start_hour += 24
            end_hour = last_interval.end_hour
            while end_hour <= start_hour:
                end_hour += 24
            start_min = (start_hour - BASE_HOUR) * 60
            end_min = (end_hour - BASE_HOUR) * 60
            duration = end_min - start_min + 60  # buffer_minutes
            timeline_meta = {
                "base_hour": BASE_HOUR,
                "duration_minutes": duration,
                "start_label": first_interval.label,
                "end_label": last_interval.label,
            }
        else:
            timeline_meta = {"base_hour": BASE_HOUR, "duration_minutes": 24 * 60, "start_label": "5-6", "end_label": "3-4"}

        return {
            "source_file": "database",
            "stations": station_names,
            "station_distances": station_distances,
            "hours": hour_labels,
            "timeline": timeline_meta,
            "directions": directions_data,
        }

    except Exception as e:
        raise DatabaseGraphError(f"Ошибка загрузки данных из базы данных: {str(e)}") from e