daterange.py

import unittest
from datetime import timedelta, date
from collections.abc import Sequence, Iterator
from itertools import islice


class date_range(Sequence):

    def __init__(self, start, stop, step=timedelta(days=1)):
        self.start = start
        self.stop = stop
        self.step = step
        self.current = None

    def __repr__(self):
        return f"{self.__class__.__name__}({self.start!r}, {self.stop!r}, step={self.step!r})"

    def __getitem__(self, index):
        return next(islice(self, index, index+1))

    def __len__(self):
        return (self.stop - self.start).days // self.step.days

    def __contains__(self, item):
        return self.start <= item < self.stop

    def __iter__(self):
        return date_range_iterator(self.start, self.stop, self.step)


class date_range_iterator(Iterator):

    def __init__(self, start, stop, step):
        self.start = start
        self.stop = stop
        self.step = step
        self.current = None

    def __next__(self):
        if self.start < self.stop and self.step < timedelta(days=0):
            raise StopIteration()
        elif self.start > self.stop and self.step > timedelta(days=0):
            raise StopIteration()
        elif self.current is None:
            self.current = self.start
        elif self.step > timedelta(days=0) and self.current + self.step >= self.stop:
            raise StopIteration()
        elif self.step < timedelta(days=0) and self.current + self.step <= self.stop:
            raise StopIteration()
        else:
            self.current = self.current + self.step
        return self.current



class DateRangeTestCase(unittest.TestCase):

    def test_init(self):
        dr = date_range(date(2024, 1, 1), date(2024, 12, 31))

    def test_repr(self):
        dr = date_range(date(2024, 1, 1), date(2024, 12, 31))
        self.assertEqual(repr(dr), "date_range(datetime.date(2024, 1, 1), datetime.date(2024, 12, 31), step=datetime.timedelta(days=1))")

    def test_getitem(self):
        dr = date_range(date(2024, 1, 1), date(2024, 12, 31))
        self.assertEqual(dr[0], date(2024, 1, 1))
        self.assertEqual(dr[9], date(2024, 1, 10))

    def test_contains(self):
        dr = date_range(date(2024, 1, 1), date(2024, 12, 31))
        self.assertIn(date(2024, 2, 2), dr)
        self.assertNotIn(date(2023, 2, 2), dr)
        self.assertIn(date(2024, 12, 30), dr)
        self.assertNotIn(date(2024, 12, 31), dr)

    def test_len(self):
        dr = date_range(date(2024, 2, 1), date(2024, 3, 1))
        self.assertEqual(len(dr), 29)
        dr = date_range(date(2024, 1, 1), date(2024, 1, 11))
        self.assertEqual(len(dr), 10)
        dr = date_range(date(2024, 1, 1), date(2024, 1, 11), step=timedelta(days=2))
        self.assertEqual(len(dr), 5)

    def test_negative_step(self):
        dr = date_range(date(2024, 3, 1), date(2024, 2, 1), step=timedelta(days=-1))
        di = iter(dr)
        self.assertEqual(next(di), date(2024, 3, 1))
        self.assertEqual(next(di), date(2024, 2, 29))

    def test_immediate_stop(self):
        dr = date_range(date(2024, 2, 1), date(2024, 3, 1), step=timedelta(days=-1))
        di = iter(dr)
        with self.assertRaises(StopIteration):
            next(di)
        dr = date_range(date(2024, 3, 1), date(2024, 2, 1), step=timedelta(days=1))
        di = iter(dr)
        with self.assertRaises(StopIteration):
            next(di)

    def test_iter(self):
        dr = date_range(date(2024, 6, 1), date(2024, 6, 5))
        expected = [
            date(2024, 6, 1),
            date(2024, 6, 2),
            date(2024, 6, 3),
            date(2024, 6, 4),
        ]
        self.assertEqual(list(dr), expected)