Python实现类似range/linspace的灵活datetime范围生成函数

原代码问题汇总 #

• 变量名错误：step计算分支误用未定义的end变量，实际入参名为stop
• 类型不兼容：除法得到的div为浮点类型，无法直接传入range()函数
• 边界逻辑错误：step模式下的区间划分逻辑不符合要求，容易出现元素多算/漏算
• 分支冗余：手动枚举四种首尾组合计算div的写法可维护性极低

优化实现代码 #

import datetime as dt
from typing import Generator

def datetime_range(
    start: dt.datetime, 
    stop: dt.datetime, 
    step: dt.timedelta = None, 
    num: int = None, 
    startpoint: bool = True, 
    endpoint: bool = True
) -> Generator[dt.datetime, None, None]:
    # 入参基础校验
    assert bool(step) != bool(num), "仅允许传入step、num其中一个参数"
    assert stop > start, "结束时间必须晚于开始时间"
    delta = stop - start

    # 固定生成数量的num模式
    if num is not None:
        assert num > 0, "num必须为正整数"
        # 计算区间总间隔数
        if startpoint and endpoint:
            interval_cnt = num - 1
        elif startpoint != endpoint:
            interval_cnt = num
        else:
            interval_cnt = num + 1
        # 兼容num=1的特殊场景
        if interval_cnt == 0:
            yield start
            return
        step = delta / interval_cnt
        # 计算起始偏移位置
        current = start if startpoint else start + step
        # 生成指定数量的时间序列
        for _ in range(num):
            yield current
            current += step
        return

    # 固定步长的step模式
    current = start if startpoint else start + step
    while True:
        if (endpoint and current > stop) or (not endpoint and current >= stop):
            break
        yield current
        current += step

效果验证 #

使用你提供的测试用例运行：

start = dt.datetime.min
stop = start + dt.timedelta(days=1)
num = 4

# 首尾都包含
print(list(datetime_range(start, stop, num=num, startpoint=True, endpoint=True)))
# 输出：[datetime.datetime(1, 1, 1, 0, 0), datetime.datetime(1, 1, 1, 8, 0), datetime.datetime(1, 1, 1, 16, 0), datetime.datetime(1, 1, 2, 0, 0)]

# 不包含起点包含终点
print(list(datetime_range(start, stop, num=num, startpoint=False, endpoint=True)))
# 输出：[datetime.datetime(1, 1, 1, 6, 0), datetime.datetime(1, 1, 1, 12, 0), datetime.datetime(1, 1, 1, 18, 0), datetime.datetime(1, 1, 2, 0, 0)]

# 包含起点不包含终点
print(list(datetime_range(start, stop, num=num, startpoint=True, endpoint=False)))
# 输出：[datetime.datetime(1, 1, 1, 0, 0), datetime.datetime(1, 1, 1, 6, 0), datetime.datetime(1, 1, 1, 12, 0), datetime.datetime(1, 1, 1, 18, 0)]

# 首尾都不包含
print(list(datetime_range(start, stop, num=num, startpoint=False, endpoint=False)))
# 输出：[datetime.datetime(1, 1, 1, 4, 48), datetime.datetime(1, 1, 1, 9, 36), datetime.datetime(1, 1, 1, 14, 24), datetime.datetime(1, 1, 1, 19, 12)]

完全符合预期输出要求，且step模式下通过循环判断边界避免了浮点除法的精度问题，逻辑更简洁稳定。

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