用于灵活创建日期时间范围的 Python 函数

Question

我正在尝试编写一个可以灵活创建连续日期时间对象的函数，类似于标准range() 函数或numpy linspace() 函数。 该函数应该能够接受num 参数或step 大小参数来确定序列。此外，应该有包括或排除起点和终点的选项。 因此，函数的签名应该是这样的：

import datetime as dt
def datetime_range(start: dt.datetime, stop: dt.datetime, step: dt.timedelta = None, num: int = None, startpoint=True, endpoint=True)

所需的输出应如下所示：

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

>>> datetime_range(start, end, 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)]

>>> datetime_range(start, end, 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)]

>>> datetime_range(start, end, 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)]

>>> datetime_range(start, end, 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)]

到目前为止，我想出的内容如下。但它没有按预期工作，我被卡住了。

def datetime_range(start: dt.datetime, stop: dt.datetime, step: dt.timedelta = None, num: int = None, startpoint=True,endpoint=True):

    assert bool(step) != bool(num), f'only one of step or num must be given'
    delta = stop - start
    if bool(num):
        if endpoint and startpoint:
            div = num - 1
        elif endpoint != startpoint:
            div = num
        else:
            div = num + 1
        step = delta / div
    else:
        if endpoint and startpoint:
            div = delta / step
        elif endpoint != startpoint:
            div = (end - step) / step  # or (delta / step) - 1
        else:
            div = (end - 2 * step) / step  # or (delta/step) -2

    return ((start + (not startpoint) * step) + x * step for x in range(div))

我正在使用 Python 3.8。也许有人可以给我提示。另外，我怀疑计算div 的方法可能比使用if, elif, else 更优雅。非常感谢任何帮助！

Answer 1

更正您的代码：

• 使用num时，返回的项目的长度必须为num，例如如果 num=4 则预期结果列表包含 4 个日期时间。所以.. for x in range(div) 是不正确的，因为div 可能是+- num，因此您的代码在预期为4 时仅返回3 项的主要原因。
• 使用step 时，我们不必担心迭代的总数，因为递增的step 是固定的。我们唯一需要做的就是将这些增量作为range 的基础，同时考虑是否包含startpoint 和endpoint。
• 您应该进一步验证输入，以确保 step 或 num 不为 0，以确保我们正在前进到 stop 计数，并且不会因为仅增加 0 而停留在同一点。
• 最好在生成器外部调整start。这样更直接。
• bool(num) 等使用的条件不是必需的。验证后直接查看num即可。

import datetime as dt


def datetime_range(start: dt.datetime, stop: dt.datetime, step: dt.timedelta = None, num: int = None, startpoint=True,endpoint=True):
    assert None in (step, num), f'only one of step or num must be given'
    if step is not None:
        assert step.total_seconds() > 0, f"Step shouldn't be 0"
    if num is not None:
        assert num > 0, f"Num shouldn't be 0"

    delta = stop - start

    if num:
        if endpoint and startpoint:
            div = num - 1
        elif endpoint or startpoint:
            div = num
        else:
            div = num + 1
        step = delta / div
    else:
        num = delta // step
        if endpoint and startpoint:
            num += 1
        elif endpoint or startpoint:
            pass
        else:
            num -= 1

    if not startpoint:
        start += step

    return (start + (x * step) for x in range(num))


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

print(f"Using num {num}")
print(list(datetime_range(start, stop, num=num, startpoint=True, endpoint=True)))
print(list(datetime_range(start, stop, num=num, startpoint=False, endpoint=True)))
print(list(datetime_range(start, stop, num=num, startpoint=True, endpoint=False)))
print(list(datetime_range(start, stop, num=num, startpoint=False, endpoint=False)))

print(f"\nUsing step {step}")
print(list(datetime_range(start, stop, step=step, startpoint=True, endpoint=True)))
print(list(datetime_range(start, stop, step=step, startpoint=False, endpoint=True)))
print(list(datetime_range(start, stop, step=step, startpoint=True, endpoint=False)))
print(list(datetime_range(start, stop, step=step, startpoint=False, endpoint=False)))

输出

Using num 4
[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)]
[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)]
[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)]
[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)]

Using step 4:00:00
[datetime.datetime(1, 1, 1, 0, 0), datetime.datetime(1, 1, 1, 4, 0), datetime.datetime(1, 1, 1, 8, 0), datetime.datetime(1, 1, 1, 12, 0), datetime.datetime(1, 1, 1, 16, 0), datetime.datetime(1, 1, 1, 20, 0), datetime.datetime(1, 1, 2, 0, 0)]
[datetime.datetime(1, 1, 1, 4, 0), datetime.datetime(1, 1, 1, 8, 0), datetime.datetime(1, 1, 1, 12, 0), datetime.datetime(1, 1, 1, 16, 0), datetime.datetime(1, 1, 1, 20, 0), datetime.datetime(1, 1, 2, 0, 0)]
[datetime.datetime(1, 1, 1, 0, 0), datetime.datetime(1, 1, 1, 4, 0), datetime.datetime(1, 1, 1, 8, 0), datetime.datetime(1, 1, 1, 12, 0), datetime.datetime(1, 1, 1, 16, 0), datetime.datetime(1, 1, 1, 20, 0)]
[datetime.datetime(1, 1, 1, 4, 0), datetime.datetime(1, 1, 1, 8, 0), datetime.datetime(1, 1, 1, 12, 0), datetime.datetime(1, 1, 1, 16, 0), datetime.datetime(1, 1, 1, 20, 0)]