Calculate adding and subtracting dates CLI

time_math.py

#!/usr/bin/python3
from datetime import date, datetime, timedelta
import sys


def parse(s):
    """
    Parse strings in the form:
        n days from now
        - n:    integer
        - days: days | weeks
        - from: from | until | before | after
        - now:  now | today | date (mm/dd)

    TODO: days until {date}
    """
    PERIODS = set(['days', 'weeks'])
    DIR_PLUS = set(['from', 'after'])
    DIR_MINUS = set(['until', 'before'])
    DATE = set(['now', 'today'])

    s = s.lower()
    s = s.split()

    if len(s) != 4:
        raise ValueError('Invalid argument length')

    num = int(s[0])

    period = s[1]
    direction = s[2]
    day = s[3]

    if period not in PERIODS:
        raise ValueError(f'Invalid period: {period}')

    if direction in DIR_PLUS:
        diff = lambda a, b: a + b
    elif direction in DIR_MINUS:
        diff = lambda a, b: a - b
    else:
        raise ValueError(f'Invalid direction: {direction}')

    TODAY = date.today()

    try:
        day = datetime.strptime(day, '%m/%d').date()
        day = day.replace(year=TODAY.year)
    except ValueError:
        if day not in DATE:
            raise ValueError(f'Invalid date: {day}')
        day = TODAY

    result = diff(day, timedelta(**{period: num}))
    return result


if len(sys.argv) < 2:
    print('Provide date string')
    sys.exit(1)

try:
    result = parse(' '.join(sys.argv[1:]))
except ValueError as e:
    print(f'Invalid input: {e}')
    sys.exit(1)

print(result)