src.fairreckitlib.core.pipeline.core_pipeline

  1"""This module contains the base core pipeline class.
  2
  3Classes:
  4
  5    CorePipeline: base pipeline class with shared IO functionality.
  6
  7This program has been developed by students from the bachelor Computer Science at
  8Utrecht University within the Software Project course.
  9© Copyright Utrecht University (Department of Information and Computing Sciences)
 10"""
 11
 12import time
 13
 14import pandas as pd
 15
 16from ..events.event_dispatcher import EventDispatcher
 17from ..events.event_error import ON_RAISE_ERROR, ErrorEventArgs
 18from ..io.event_io import ON_CREATE_FILE, DataframeEventArgs, FileEventArgs
 19
 20
 21class CorePipeline:
 22    """Base class for FairRecKit pipelines.
 23
 24    This class exposes some reusable functionality that can be used in derived
 25    pipelines to read and/or write dataframes.
 26    """
 27
 28    def __init__(self, event_dispatcher: EventDispatcher):
 29        """Construct the CorePipeline.
 30
 31        Args:
 32            event_dispatcher: used to dispatch events when running the core pipeline.
 33        """
 34        self.event_dispatcher = event_dispatcher
 35
 36    def read_dataframe(
 37            self,
 38            dataframe_path: str,
 39            dataframe_name: str,
 40            event_id_on_begin: str,
 41            event_id_on_end: str,
 42            *,
 43            names=None) -> pd.DataFrame:
 44        """Read a dataframe from the disk.
 45
 46        This function dispatches an error event when the FileNotFoundError is raised,
 47        and thereafter the error is raised once more.
 48
 49        Args:
 50            dataframe_path: path to the dataframe file.
 51            dataframe_name: name of the dataframe to use for event dispatching.
 52            event_id_on_begin: the event_id to dispatch when loading starts.
 53            event_id_on_end: the event_id to dispatch when loading is finished.
 54            names: the column names of the dataframe or None to infer them from the header.
 55
 56        Raises:
 57            FileNotFoundError: when the dataframe file is not found.
 58
 59        Returns:
 60            the loaded dataframe.
 61        """
 62        self.event_dispatcher.dispatch(DataframeEventArgs(
 63            event_id_on_begin,
 64            dataframe_path,
 65            dataframe_name
 66        ))
 67
 68        start = time.time()
 69
 70        try:
 71            dataframe = pd.read_csv(
 72                dataframe_path,
 73                sep='\t',
 74                header='infer' if names is None else None,
 75                names=names
 76            )
 77        except FileNotFoundError as err:
 78            self.event_dispatcher.dispatch(ErrorEventArgs(
 79                ON_RAISE_ERROR,
 80                'FileNotFoundError: raised while trying to load the ' +
 81                dataframe_name + ' from ' + dataframe_path
 82            ))
 83            raise err
 84
 85        end = time.time()
 86
 87        self.event_dispatcher.dispatch(DataframeEventArgs(
 88            event_id_on_end,
 89            dataframe_path,
 90            dataframe_name
 91        ), elapsed_time=end - start)
 92
 93        return dataframe
 94
 95    def write_dataframe(
 96            self,
 97            dataframe_path: str,
 98            dataframe: pd.DataFrame,
 99            header: bool) -> None:
100        """Write a dataframe to the disk.
101
102        This function is intended to write (append) a dataframe in chunks,
103        including the header of the dataframe.
104        It is assumed that when the header is True that the dataframe file
105        has been created, which in turn will dispatch the IO event.
106
107        Args:
108            dataframe_path: path to the dataframe file.
109            dataframe: the dataframe to append to the file.
110            header: whether to include the header.
111        """
112        dataframe.to_csv(
113            dataframe_path,
114            mode='a',
115            sep='\t',
116            header=header,
117            index=False
118        )
119        # header is the first line meaning the file has just been created
120        if header:
121            self.event_dispatcher.dispatch(FileEventArgs(ON_CREATE_FILE, dataframe_path))