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),
)
)