def test_dataframe_strategy(pdtype, data):
"""Test DataFrameSchema strategy."""
dataframe_schema = pa.DataFrameSchema(
{f"{pdtype.value}_col": pa.Column(pdtype)}
)
dataframe_schema(data.draw(dataframe_schema.strategy(size=5)))
with pytest.raises(pa.errors.BaseStrategyOnlyError):
strategies.dataframe_strategy(
pdtype, strategies.pandas_dtype_strategy(pdtype)
)
def test_dataframe_strategy(pdtype, size, data):
"""Test DataFrameSchema strategy."""
dataframe_schema = pa.DataFrameSchema(
{f"{pdtype.value}_col": pa.Column(pdtype)})
df_sample = data.draw(dataframe_schema.strategy(size=size))
if size == 0:
assert df_sample.empty
elif size is None:
assert df_sample.empty or isinstance(dataframe_schema(df_sample),
pd.DataFrame)
else:
assert isinstance(dataframe_schema(df_sample), pd.DataFrame)
with pytest.raises(pa.errors.BaseStrategyOnlyError):
strategies.dataframe_strategy(pdtype,
strategies.pandas_dtype_strategy(pdtype))
def test_check_nullable_dataframe_strategy(data_type, nullable, data):
"""Test strategies for generating nullable DataFrame data."""
size = 5
# pylint: disable=no-value-for-parameter
strat = strategies.dataframe_strategy(
columns={"col": pa.Column(data_type, nullable=nullable, name="col")},
size=size,
)
example = data.draw(strat)
if nullable:
assert example.isna().sum(axis=None).item() >= 0
else:
assert example.notna().all(axis=None)
def test_check_nullable_dataframe_strategy(pdtype, nullable, data):
"""Test strategies for generating nullable DataFrame data."""
size = 5
strat = strategies.dataframe_strategy(
columns={
"col": pa.Column(
pandas_dtype=pdtype, nullable=nullable, name="col"
)
},
size=size,
)
example = data.draw(strat)
if nullable:
assert example.isna().any(axis=None)
else:
assert example.notna().all(axis=None)