如何用Hypothesis生成非空且无全NaN列的pd.DataFrame

要解决你遇到的两个问题——避免空DataFrame、确保没有全NaN的列，只需对原代码做两处关键修改：

1. 强制生成非空索引 #

给indexes添加min_size=1参数，直接限制索引至少包含一行数据，从根源上杜绝空DataFrame的生成。

2. 过滤全NaN列的情况 #

利用Hypothesis的st.data()工具生成DataFrame后，添加过滤条件，确保每一列都存在至少一个非NaN值。

修改后的完整代码 #

import datetime
import pandas as pd
from hypothesis import strategies as st, data
from hypothesis.extra.pandas import columns as cols, data_frames, indexes

@data
def generate_valid_df(d):
    # 生成带非空索引的DataFrame
    df = d.draw(data_frames(
        columns=cols(
            ["sec1", "sec2", "sec3"], 
            elements=st.floats(allow_infinity=False)
        ),
        index=indexes(
            elements=st.dates(
                min_value=datetime.date(2023,10,31), 
                max_value=datetime.date(2024,5,31)
            ),
            min_size=1  # 确保至少一行数据
        )
    ))
    # 过滤掉存在全NaN列的结果
    for col in df.columns:
        if df[col].isna().all():
            raise ValueError(f"列 {col} 全为NaN，重新生成")
    return df

# 生成符合要求的示例
print(generate_valid_df.example())

另一种更简洁的写法 #

如果不想用装饰器，也可以直接用filter方法过滤不符合条件的DataFrame：

import datetime
import pandas as pd
from hypothesis import strategies as st
from hypothesis.extra.pandas import columns as cols, data_frames, indexes

valid_df_strategy = st.data().flatmap(lambda d:
    d.draw(data_frames(
        columns=cols(
            ["sec1", "sec2", "sec3"], 
            elements=st.floats(allow_infinity=False)
        ),
        index=indexes(
            elements=st.dates(
                min_value=datetime.date(2023,10,31), 
                max_value=datetime.date(2024,5,31)
            ),
            min_size=1
        )
    )).filter(lambda df: all(not df[col].isna().all() for col in df.columns))
)

print(valid_df_strategy.example())

这两种写法都能保证：

• 生成的DataFrame至少有一行数据
• 每一列都包含至少一个非NaN的浮点值

内容的提问来源于stack exchange，提问作者Andi