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!