DataImporter.py

import fastf1
from abc import ABC

from fastf1.core import Session
from fastf1.events import EventSchedule, Event

from DataAnalyser import DataAnalyser
from DataHandler import DataHandler, SessionIdentifier


class DataImporter(DataHandler, ABC):
    '''
    Imports/loads sessions as specified by SessionIdentifiers and/or certain criteria such as rain.

    Any method of this class must be given a SessionIdentifier, a time or a time period by which to select
    sessions. If the method does not require one of these, it should not be part of this class.
    '''


    def importAllEventsFromYear(self, year: int):
        races: list[Event] = []
        schedule: EventSchedule = fastf1.get_event_schedule(year, include_testing = False)
        for raceIndex in schedule['RoundNumber']:
            races.append(schedule.get_event_by_round(raceIndex))
        return races

    def importSessionWeather(self, sessionIdentifier: SessionIdentifier):
        '''
        Import a session containing only weather data.
        :param sessionIdentifier: Session to import. Note that only races after 2017 can import weather, as session data
        before 2018 is provided by the Ergast/Jolpica API, which does not contain weather data.
        :return: Session object containing only weather data and basic session information.
        '''
        return self.importSession(sessionIdentifier, laps = False, telemetry = False, weather = True, messages = False)

    def importSessions(self, sessionIdentifiers: list[SessionIdentifier], laps = True, telemetry = False, weather = False, messages = False):
        sessions: list[Session] = []
        for sessionIdentifier in sessionIdentifiers:
            sessions.append(self.importSession(sessionIdentifier, laps, weather, messages, telemetry))
        return sessions

    def importSession(self, sessionIdentifier: SessionIdentifier, laps = True, telemetry = False, weather = False, messages = False):
        if weather is True and sessionIdentifier.year < self.firstFastF1Year:
            raise ValueError(f"Cannot fetch weather data for races before {self.firstFastF1Year} due to API limitations")

        session = fastf1.get_session(sessionIdentifier.year, sessionIdentifier.event, sessionIdentifier.sessionType)
        session.load(laps = laps, telemetry = telemetry, weather = weather, messages = messages)
        return session

    def getRainRacesSince(self, firstYear: int):
        rainRaces: list[Session] = []
        for firstYear in range(firstYear, 2025):  # FIXME: Add automatic upper bound for year so that it can at most be set to current year or year of latest f1 race (e.g. shortly after new years eve)
            wetWeatherRacesInYear: list[Session] = self.getRainRacesIn(firstYear)
            for wetWeatherRace in wetWeatherRacesInYear:
                rainRaces.append(wetWeatherRace)
        return rainRaces

    def getRainRacesIn(self, year: int):
        events: list[Event] = self.importAllEventsFromYear(year)
        races: list[Session] = []
        for event in events:
            raceIdentifier: SessionIdentifier = SessionIdentifier(event.year, event["RoundNumber"], "R")
            try:
                raceSession: Session = self.importSessionWeather(raceIdentifier)
                races.append(raceSession)
            except ValueError: # Needed as long as weather data for races before 2018 hasn't been supplemented
                print(f"Weather for {event["EventName"]} {event.year} could not be determined")

        analyser: DataAnalyser = DataAnalyser()
        rainRaces: list[Session] = analyser.filterForWetSessions(races)

        return rainRaces