from classes.data_operator import (
    DataFrameOperator,
    ExtendedDataFrame,
    ExtendedDataFrameList,
)
from scheduler.block import BaseBlock


class DataFrameProcessing(BaseBlock):
    def __init__(self, operator: DataFrameOperator):
        self.operator = operator

    def __call__(self, input: ExtendedDataFrame, **kwargs) -> ExtendedDataFrame:
        response = self.operator(input, **kwargs)
        if hasattr(self.operator, "usage"):
            self.usage = self.operator.usage
        return response


class BatchDataFrameProcessing(BaseBlock):
    def __init__(self, operator: DataFrameOperator):
        self.operator = operator

    def __call__(self, input: ExtendedDataFrameList, **kwargs) -> list:
        return [self.operator(df, **kwargs) for df in input.list_of_df]