如何用Hypothesis为测试DataFrame创建pd.datetime类型列？

我想要测试某一方法能否处理作为参数传入的pandas DataFrame中的不同日期，但使用column('Date', dtype=pd.datetime)无法在测试DataFrame中创建日期列，示例代码如下：

from hypothesis import given
from hypothesis.extra.pandas import column, data_frames
import pandas as pd
from unittest import TestCase

class TestExampleClass(TestCase):

    @given(data_frames([column('A', dtype=str), column('B', dtype=int),column('Date', dtype=pd.datetime)]))
    def test_example_test_method(self, dataframe):
        self.assertTrue(True)

我了解相关的datetime索引DataFrame创建方法，但无法解决当前场景的问题。

问题出在pd.datetime不是Hypothesis column函数支持的有效dtype参数（当前pandas已推荐使用pd.Timestamp或pd.DatetimeDtype替代），要生成合法的日期列，需要用Hypothesis的日期生成策略来指定列的数据规则，而非直接传递dtype。

方案1：用elements策略生成Timestamp列 #

直接给column指定elements参数，借助hypothesis.strategies.datetimes()生成datetime对象，pandas会自动将该列识别为datetime类型：

from hypothesis import given
from hypothesis.extra.pandas import column, data_frames
from hypothesis.strategies import datetimes
import pandas as pd
from unittest import TestCase

class TestExampleClass(TestCase):

    @given(data_frames([
        column('A', dtype=str),
        column('B', dtype=int),
        column('Date', elements=datetimes())
    ]))
    def test_example_test_method(self, dataframe):
        # 可选：验证Date列的类型是否正确
        self.assertTrue(pd.api.types.is_datetime64_any_dtype(dataframe['Date']))
        self.assertTrue(True)

方案2：明确指定pd.DatetimeDtype #

如果需要显式声明列的dtype，可以使用pd.DatetimeDtype()，同时配合datetimes()策略保证生成的数据符合类型要求：

from hypothesis import given
from hypothesis.extra.pandas import column, data_frames
from hypothesis.strategies import datetimes
import pandas as pd
from unittest import TestCase

class TestExampleClass(TestCase):

    @given(data_frames([
        column('A', dtype=str),
        column('B', dtype=int),
        column('Date', dtype=pd.DatetimeDtype(), elements=datetimes())
    ]))
    def test_example_test_method(self, dataframe):
        self.assertTrue(pd.api.types.is_datetime64_any_dtype(dataframe['Date']))
        self.assertTrue(True)

额外优化 #

你可以通过datetimes()的参数限制日期范围，生成更贴合测试场景的数据，比如：

datetimes(min_value=pd.Timestamp('2020-01-01'), max_value=pd.Timestamp('2024-12-31'))

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