Skip to content
Advertisement

Is there a way to specify multiple columns of a dataframe with different strategies in Hypothesis?

I would like to create a dataframe for testing that looks like:

target var0 var1 var2 var3
0 0.34 1.43 0.41 0.98
1 -1.43 -0.31 7.43 1.95

I have been able to do this by defining the columns as a dictionary as seen in this answer https://stackoverflow.com/a/65720110/251754, but I’m wondering if there is a better way to do it. I have tried the following, which doesn’t work because elements cannot take a list:

@given(x_df=
    data_frames(
            columns=columns(names_or_number=['target','var0','var1','var2','var3'],
            elements=[strategies.integers(), strategies.floats()]),
        index=range_indexes(min_size=10, max_size=10),
    )
)

Advertisement

Answer

You can build up a list of column objects by hand, e.g. by:

@given(
    x_df=data_frames(
        columns=[column(name="target", elements=strategies.integers())] + columns(
            names_or_number=["var0", "var1", "var2", "var3"],
            elements=strategies.floats(),
        ),
        index=range_indexes(min_size=10, max_size=10),
    )
)
Advertisement