DataImporter.py

import re
import string
from datetime import datetime

import fastf1
from abc import ABC

import requests
from fastf1.core import Session
from fastf1.ergast import fetch_day
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 public method of this class must be given an Event, Session, SessionIdentifier, time or 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 fetchWeather(self, event: Event | SessionIdentifier | list[Event] | list[SessionIdentifier]):
        """
        Fetch session weather from the wikipedia entry for the specified event(s).
        :param event: The event or identifier for any session of the event to fetch weather for. Alternatively a list of
        these.
        :return: Weather conditions as string, or a list of such strings if a list is provided.
        """

        # Recursive call if list provided
        if isinstance(event, list):
            weatherEntries: list[string] = []
            for e in event:
                weatherEntries.append(self.fetchWeather(e))
            return weatherEntries

        # Convert from SessionIdentifier to Event if needed
        if isinstance(event, SessionIdentifier):
            event = self.importEvent(event.year, event.event)

        # Fetch weather
        wikiHtml = self.__importWikiHtml(event)
        analyser = DataAnalyser()
        weather = analyser.extractWeatherFromHtml(wikiHtml)

        return weather

    def __importWikiHtml(self, event: Event):
        """
        Fetch the HTML contents of the wikipedia page for the event given.
        :param event: Event whose wikipedia page to fetch.
        :return: HTML content of the wikipedia page as a string.
        """
        apiRootUrl: string = "https://en.wikipedia.org/wiki/"
        grandPrixName: string = f"{event.year} {event["EventName"]}"
        uri: string = re.sub(" ", "_", grandPrixName)
        url: string = f"{apiRootUrl}{uri}"
        response = requests.get(url)
        if not response.ok:
            raise ConnectionError(f"Could not fetch wikipedia article with URL {url}")
        return response.text
        # URL example: https://en.wikipedia.org/w/api.php?action=query&titles=2024_S%C3%A3o_Paulo_Grand_Prix&prop=extracts&format=json&exintro=1

    def importEvent(self, year: int, event: int | str):
        return fastf1.get_event(year, event)

    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")

        if sessionIdentifier.year < self.firstFastF1Year:
            print(f"WARNING: No FastF1 data available before {self.firstFastF1Year}, using Jolpica data instead. Some expected data may be missing")

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

    def getWetEventsSince(self, firstYear: int):
        currentYear = datetime.now().year
        if firstYear > currentYear:
            raise ValueError("Cannot get race data from the future :)")

        rainRaces: list[Session] = []
        for firstYear in range(firstYear, currentYear):  # FIXME: Handle exception after new years, when no events have run in current year yet
            wetWeatherRacesInYear: list[Session] = self.getWetEventsIn(firstYear)
            for wetWeatherRace in wetWeatherRacesInYear:
                rainRaces.append(wetWeatherRace)
        return rainRaces

    # TODO: Rename (other than legacy)
    @DeprecationWarning
    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.filterForRainSessionsLEGACY(races)

        return rainRaces

    def getWetEventsIn(self, year: int):
        print(f"Fetching wet events in {year}")

        events: list[Event] = self.importAllEventsFromYear(year)
        analyser: DataAnalyser = DataAnalyser()
        rainRaces: list[Session] = analyser.filterForWetEvents(events)

        return rainRaces