3
\$\begingroup\$

I've created a function to split a period into the biggest different time unit possible. It can be years, months, weeks or days.
However, it ended up being quite big and I feel it can be greatly improved, but I'm having difficulties finding ways to do it.

import calendar
from datetime import datetime, timedelta


def create_intervals(start_date, end_date):
    start = datetime.strptime(start_date, "%Y-%m-%d")
    end = datetime.strptime(end_date, "%Y-%m-%d")
    di = []
    # Year
    if (sy := start.year) < (ey := end.year):
        for y in range(sy, ey + 1):
            if sy == y:
                di.append((start_date, f"{y}-12-31"))
            elif sy < y < ey:
                di.append((f"{y}-01-01", f"{y}-12-31"))
            else:
                di.append((f"{y}-01-01", end_date))
    # Month
    elif (sm := start.month) < (em := end.month):
        for m in range(sm, em + 1):
            last_day = calendar.monthrange(sy, m)[1]
            if sm == m:
                di.append((start_date, f"{sy}-{m:02d}-{last_day}"))
            elif sm < m < em:
                di.append((f"{sy}-{m:02d}-01", f"{sy}-{m:02d}-{last_day}"))
            else:
                di.append((f"{sy}-{m:02d}-01", end_date))
    # Week
    elif (sw := start.isocalendar().week) < (ew := end.isocalendar().week):
        for w in range(sw, ew + 1):
            if sw == w:
                last = start - timedelta(days=start.weekday()) + timedelta(days=6)
                di.append((start_date, f"{sy}-{sm:02d}-{last.day:02d}"))
            elif sw < w < ew:
                wd = datetime.strptime(f"{sy}-W{w}-1", "%Y-W%W-%w")
                first = wd - timedelta(days=wd.weekday())
                last = wd - timedelta(days=wd.weekday()) + timedelta(days=6)
                di.append((f"{sy}-{sm:02d}-{first.day:02d}", f"{sy}-{sm:02d}-{last.day:02d}"))
            else:
                first = end - timedelta(days=end.weekday())
                di.append((f"{sy}-{sm:02d}-{first.day:02d}", end_date))
    # Day
    elif start.day < end.day:
        for d in range(start.day, end.day + 1):
            di.append((f"{sy}-{sm:02d}-{d:02d}", f"{sy}-{sm:02d}-{d:02d}"))
    # None
    else:
        di.append((start_date, end_date))
    return di

Example for years as biggest different time unit:

create_intervals("2021-10-22", "2023-03-02")
# Outputs intervals for each year:
# [("2021-10-22", "2021-12-31"),
#  ("2022-01-01", "2022-12-31"),
#  ("2023-01-01", "2023-03-02")]

Example for days as biggest different time unit:

create_intervals("2021-10-22", "2021-10-24")
# Outputs:
# [("2021-10-22", "2021-10-22"),
#  ("2021-10-23", "2021-10-23"),
#  ("2021-10-24", "2021-10-24")]
\$\endgroup\$
1
  • \$\begingroup\$ Could you also show an invocation example for weeks? \$\endgroup\$
    – Reinderien
    Oct 26, 2022 at 23:36

1 Answer 1

4
\$\begingroup\$

There is never a case where you should use string datetime representation in this code, and there is never a case where you should use datetime; use date instead.

Consider converting to a generator function to simplify your code.

Your one- and two-letter variable names need to go away.

Almost all of your walrus assignments need to go away except that of the week comparison.

first = wd - timedelta(days=wd.weekday()) should not do any subtraction because the right-hand term will always be 0.

Write unit tests based on the cases you already have, and expand for cases that you don't already have.

Writing conditions for the beginning and end of a sequence within a loop is an antipattern - instead, cut those out to statements before and after the loop, and make the loop unconditional.

Add PEP484 typehints.

Suggested

import calendar
from datetime import date, timedelta
from typing import Iterator


def create_intervals(start: date, end: date) -> Iterator[tuple[date, date]]:
    # Year
    if start.year < end.year:
        yield start, date(start.year, 12, 31)
        for year in range(start.year + 1, end.year):
            yield date(year, 1, 1), date(year, 12, 31)
        yield date(end.year, 1, 1), end

    # Month
    elif start.month < end.month:
        last_days = (
            calendar.monthrange(start.year, month)[1]
            for month in range(start.month, end.month + 1)
        )

        yield start, date(start.year, start.month, next(last_days))

        for month, last_day in zip(range(start.month + 1, end.month), last_days):
            yield date(start.year, month, 1), date(start.year, month, last_day)

        yield date(start.year, end.month, 1), end

    # Week
    elif (start_week := start.isocalendar().week) < (end_week := end.isocalendar().week):
        last = start + timedelta(days=6 - start.weekday())
        yield start, date(start.year, start.month, last.day)

        for week in range(start_week + 1, end_week):
            first = date.fromisocalendar(start.year, week, 1)
            last = first + timedelta(days=6)
            yield date(start.year, start.month, first.day), date(start.year, start.month, last.day)

        first = end - timedelta(days=end.weekday())
        yield date(start.year, start.month, first.day), end

    # Day
    elif start.day < end.day:
        for day in range(start.day, end.day + 1):
            yield date(start.year, start.month, day), date(start.year, start.month, day)

    # None
    else:
        yield start, end


def test() -> None:
    # Years
    result = tuple(create_intervals(date(2021, 10, 22), date(2023,  3,  2)))
    assert result == (
        (date(2021, 10, 22), date(2021, 12, 31)),
        (date(2022,  1,  1), date(2022, 12, 31)),
        (date(2023,  1,  1), date(2023,  3,  2)),
    )

    # Months
    result = tuple(create_intervals(date(2021, 3, 7), date(2021, 10, 26)))
    assert result == (
        (date(2021,  3, 7), date(2021,  3, 31)),
        (date(2021,  4, 1), date(2021,  4, 30)),
        (date(2021,  5, 1), date(2021,  5, 31)),
        (date(2021,  6, 1), date(2021,  6, 30)),
        (date(2021,  7, 1), date(2021,  7, 31)),
        (date(2021,  8, 1), date(2021,  8, 31)),
        (date(2021,  9, 1), date(2021,  9, 30)),
        (date(2021, 10, 1), date(2021, 10, 26)),
    )

    # Weeks
    result = tuple(create_intervals(date(2021, 10, 2), date(2021, 10, 26)))
    assert result == (
        (date(2021, 10,  2), date(2021, 10,  3)),
        (date(2021, 10,  4), date(2021, 10, 10)),
        (date(2021, 10, 11), date(2021, 10, 17)),
        (date(2021, 10, 18), date(2021, 10, 24)),
        (date(2021, 10, 25), date(2021, 10, 26))
    )

    # Days
    result = tuple(create_intervals(date(2021, 10, 22), date(2021, 10, 24)))
    assert result == (
        (date(2021, 10, 22), date(2021, 10, 22)),
        (date(2021, 10, 23), date(2021, 10, 23)),
        (date(2021, 10, 24), date(2021, 10, 24)),
    )


if __name__ == '__main__':
    test()
\$\endgroup\$
1
  • 1
    \$\begingroup\$ This was such an improvement! I really liked how converting to a generator function and using the unconditional loops simplified the whole structure. Thank you for taking the time \$\endgroup\$
    – CIRCLE
    Oct 28, 2022 at 1:13

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.