POC descriptor tests

[mike0sv]

from typing import Union, Dict

import pandas as pd

from evidently import Dataset, ColumnType
from evidently.core.datasets import Descriptor, DatasetColumn
from evidently.legacy.base_metric import DisplayName
from evidently.legacy.options.base import Options
from evidently.tests import eq


class Add1(Descriptor):
    column: str

    def generate_data(self, dataset: "Dataset", options: Options) -> Union[
        DatasetColumn, Dict[DisplayName, DatasetColumn]]:
        col = dataset.column(self.column)
        return DatasetColumn(type=ColumnType.Numerical, data=col.data + 1)


class Add12(Descriptor):
    column: str

    def generate_data(self, dataset: "Dataset", options: Options) -> Union[
        DatasetColumn, Dict[DisplayName, DatasetColumn]]:
        col = dataset.column(self.column)
        return {
            "add 1": DatasetColumn(type=ColumnType.Numerical, data=col.data + 1),
            "add 2": DatasetColumn(type=ColumnType.Numerical, data=col.data + 2),
        }


def main():
    data = {"col": [1, 2, 3]}
    dataset = Dataset.from_pandas(
        pd.DataFrame(data),
        descriptors=[
            Add1(column="col", alias="col_plus_1", tests=[eq(2)]),
            Add12(column="col", alias="col_plus_12", tests=[eq(3)]),
            # aggregated score/test:
            # CaseEval(success=True, rate=True, score=True, all=True, any=True,
            #          score_weights={"summary_negativity": 0.9, "summary_length": 0.1},
            #          alias="case_eval"),
        ],
    )
    df = dataset.as_dataframe()
    for col in df.columns:
        print(col)
        print(df[col].tolist())
        print()


if __name__ == '__main__':
    main()