Python function for flexibly creating datetime range-CodePudding

I am trying to write a function that can flexibly create consecutive datetime objects similarly to the standard range() function or the numpy linspace() function. The function should be able to accept either a num argument or a step size argument to determine the sequence. Moreover, there should be options to include or exclude both the startpoint and the endpoint. Thus, the signature of the function should be like this:

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

Desired outputs should look something like this:

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)]

What I have come up with so far is below. But it is not working as intended and I got stuck.

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))

I am working in Python 3.8. Maybe someone can give me hint. Also, I have a suspicion that there might be a more elegant way of computing the div than using if, elif, else. Any help is greatly appreciated!

CodePudding user response：

Corrections to your code:

When using num, it is mandatory that the returned items must be of length num e.g. if num=4 then it is expected that the resulting list contains 4 datetimes. So .. for x in range(div) is incorrect because div might be - num thus the main reason why your code returns only 3 items when 4 is expected.
When using step, we shouldn't worry about the total number of iterations as the incremented step is fixed. The only thing we need to do is count those increments as the basis of the range while taking into account if the startpoint and endpoint are included or not.
You should validate the inputs further to ensure that step or num isn't 0 to make sure that we are moving forward to the stop count and not stuck to the same point due to incrementing just 0.
It would be better to adjust start outside the generator. It is more straight-forward that way.
The conditions used such as bool(num) is not necessary. We can just directly check num after the validations we made.

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)))

Output

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)]