Skip to content
Snippets Groups Projects
DataClass.py 1.1 KiB
Newer Older
Rhino's avatar
Rhino committed
import pandas as pd

class MLDataFrame(pd.DataFrame):
    _metadata = ["_first_attempt", "_second_attempt", "_label"]
Rhino's avatar
Rhino committed

    @property
    def _constructor(self):
        return MLDataFrame
    def __init__(self, data, first_attempt=None, second_attempt=None, label=None, **kwargs):
Rhino's avatar
Rhino committed
        if kwargs.pop("index", None):
            raise ValueError("Cannot set a custom index on MLDataFrame.")
        super().__init__(data, **kwargs)

        self._first_attempt = first_attempt
        self._second_attempt = second_attempt
        self._label = label
Rhino's avatar
Rhino committed
    
    @property
    def first_attempt(self):
        if self._first_attempt is None:
Rhino's avatar
Rhino committed
            raise ValueError("Cannot access unspecified x_data")
        return self._first_attempt[self.index]
Rhino's avatar
Rhino committed

    @property
    def second_attempt(self):
        if self._second_attempt is None:
Rhino's avatar
Rhino committed
            raise ValueError("Cannot access unspecified y_data")
        return self._second_attempt[self.index]

    @property
    def label(self):
        if self._label is None:
            raise ValueError("Cannot access unspecified y_data")
        return self._label[self.index]