Python EU travel visa calculator

Question

Made this basic calculator for travel visa calculation since I was tired of doing it by hand multiple times and online calculators gave mixed results. It prints every day the visa is valid and displays how many days are left.

The basic rules:

• You start out with 90 days.
• After using (traveling) a day, it'll only become available again after 180 days.
• After not traveling at all for 90 days, your available days reset.

This was all we needed, there are far more complications and exceptions to these rules. Disclaimer: Nobody should use this as sole validation, always double check!

I know input handling is poor, but that isn't much of an issue for my uses. Mainly interested in improvements to the logic.

from datetime import date, timedelta

# Starting variables
start_days = 90
days_reset_delay = timedelta(days=180)
days_full_reset = 90
planned_trips = []
days_left = start_days
days_not_traveled = 0

# Handle inputs assuming proper values
visa_start_input = input("Enter visa start date (YYYY,MM,DD eg: 2020,11,6): ")
visa_start_input_split = visa_start_input.split(",")
visa_start = date(int(visa_start_input_split[0]), int(visa_start_input_split[1]), int(visa_start_input_split[2]))
visa_end = visa_start + timedelta(days=365)

print("Enter planned/completed trips. Enter \"end\" to stop.")
while True:
    trip_input = input("Enter trip start and end date (YYYY,MM,DD eg: 2020,11,8-2020,12,28): ")
    if trip_input.lower() == "end":
        break
    trip_input_split = trip_input.split("-")
    trip_start = trip_input_split[0].split(",")
    trip_end = trip_input_split[1].split(",")
    planned_trips.append([date(int(trip_start[0]), int(trip_start[1]), int(trip_start[2])),
                          date(int(trip_end[0]), int(trip_end[1]), int(trip_end[2]))])


# Day by day logic and result
for i in range((visa_end - visa_start).days + 1):
    day = visa_start + timedelta(days=i)
    days_not_traveled += 1
    for trip in planned_trips:
        if trip[0] <= day <= trip[1]:
            days_left -= 1
            days_not_traveled = 0
        if day == trip[0] + days_reset_delay:
            days_left += 1
            if trip[0] < trip[1]:
                trip[0] += timedelta(days=1)

    if days_not_traveled >= days_full_reset:
        days_left = start_days

    if day + timedelta(days=days_left) == visa_end - timedelta(days=1):
        print("VVV Optimal day to travel VVV ---------------------------------------------------------------------")

    print(day, "-- Days left:", days_left)

Answer 1

Interesting start to what seems a very useful utility. Future improvements could include console calendar-style formatting for much easier reading, or a desktop or web UI.

Your YYYY,MM,DD format is unusual. A slight adjustment to use hyphens instead of commas will bring it in line with ISO 8601 which will be more familiar, I think, to users. Or if you don't want to use that, you should use your locale's own date format, whatever that may be.

You should separate your Enter trip start and end date into two separate input calls.

Don't do this:

visa_start_input_split = visa_start_input.split(",")
visa_start = date(int(visa_start_input_split[0]), int(visa_start_input_split[1]), int(visa_start_input_split[2]))

Instead, use date.fromisoformat or datetime.strptime(...).date().

timedelta(days=365) is not accurate, and unfortunately there's no easy built-in way to do this right. relativedelta(years=1) from dateutil will do the trick.

I haven't taken the time to understand your algorithm, but mutating your trip data here:

trip[0] += timedelta(days=1)

is not a good idea, and you should find some other way that preserves the trip data as immutable. If you're able to make it immutable, then this:

for trip in planned_trips:

should become

for trip_start, trip_end in planned_trips:

One alternative algorithm is to represent the days-left adjustments as addend vectors in Numpy. This has the potential to execute more quickly, though for the scale of these data that won't matter that much. Numpy has some neat tricks like being able to construct an array of increasing dates in one call to arange. Having separated addends allows your trip data to remain immutable.

You really should have unit tests, and your unit tests should include a case that exercises the 180-day restoral logic - which your current example does not.

Example implementation

This includes:

• Unit tests for your own example
• Unit tests for an extended, multi-trip example that exercises the 180-day restoral logic
• Support for input+output date localisation
• Refactored support for output matching your current format (print_flat)
• Support for a calendar-like output (print_calendar)

from dataclasses import dataclass
from datetime import date, timedelta, datetime
from dateutil.relativedelta import relativedelta, SU
from itertools import chain
from locale import setlocale, nl_langinfo, D_FMT, LC_ALL
from typing import Optional, Tuple, Iterable, Collection, Dict
import numpy as np

# Calling setlocale(LC_ALL, '') lets it use the default locale as defined by
# the LANG variable.
setlocale(LC_ALL, '')
DATE_FMT = nl_langinfo(D_FMT)
from calendar import day_abbr, SUNDAY

START_DAYS = 90
DAYS_FULL_RESET = START_DAYS
DAYS_RESET_DELAY = 180
ONE_DAY = timedelta(days=1)

VISA_LENGTH = relativedelta(
    years=1,
    # needed if the visa is to include e.g. the day of 2021-11-06 when it had
    # started on 2020-11-06
    days=1,
)


def input_date(title: str, terminate: Optional[str] = None) -> Optional[date]:
    if terminate:
        suffix = f' or "{terminate}"'
    else:
        suffix = ''
    s = input(f'Enter {title} date ({DATE_FMT}){suffix}: ')
    if s == terminate:
        return None
    return datetime.strptime(s, DATE_FMT).date()


@dataclass(frozen=True)
class Trip:
    # As is standard for Python, this forms a half-closed interval
    start: date  # inclusive
    end: date    # exclusive

    @property
    def n_days(self) -> int:
        return (self.end - self.start).days

    def bounds(self, epoch: date) -> Tuple[
        int,  # start of decline
        int,  # end of decline
        int,  # start of include
        int,  # end of incline
    ]:
        a = (self.start - epoch).days
        b = (self.end - epoch).days
        return a, b, a + DAYS_RESET_DELAY, b + DAYS_RESET_DELAY

    @classmethod
    def from_stdin(cls) -> Iterable['Trip']:
        while True:
            inclusive_start = input_date('trip start', 'end')
            if inclusive_start is None:
                return
            # Assume that the user-provided end date is inclusive ("last day")
            exclusive_end = input_date('trip end') + ONE_DAY
            yield cls(inclusive_start, exclusive_end)

    def __lt__(self, other: 'Trip') -> bool:
        return (
            self.start < other.start
            or (self.start == other.start and self.end < other.end)
        )


@dataclass(frozen=True)
class TripDatabase:
    visa_start: date  # inclusive
    visa_end: date    # exclusive
    trips: Collection[Trip]

    @classmethod
    def from_stdin(cls) -> 'TripDatabase':
        start = input_date('visa start')
        return cls(start, start + VISA_LENGTH, sorted(Trip.from_stdin()))

    @property
    def n_days(self) -> int:
        return (self.visa_end - self.visa_start).days

    @property
    def n_trips(self) -> int:
        return len(self.trips)

    @property
    def days_left(self) -> np.ndarray:
        addends = np.zeros((self.n_trips, self.n_days), dtype=np.int16)
        in_trip = np.zeros((self.n_days + 1,), dtype=bool)
        in_trip[-1] = True  # Edge value needed to calculate last run

        for i, trip in enumerate(self.trips):
            dec_start, dec_end, inc_start, inc_end = trip.bounds(self.visa_start)
            in_trip[dec_start: dec_end] = True

            # Declining through each trip
            addends[i, dec_start: dec_end] = np.arange(-1, -1-trip.n_days, -1)
            # Constant days-used after the trip ...
            addends[i, dec_end: inc_start] = -trip.n_days
            # ... until the recovery period, inclining back to 0
            addends[i, inc_start: inc_end] = np.arange(-trip.n_days, 0, 1)

        # Find all of the indices where a reset occurs
        trip_idx = np.flatnonzero(in_trip)
        trip_runs = np.diff(trip_idx) - 1
        # We're too late
        resets = trip_idx[:-1][trip_runs > DAYS_FULL_RESET] + DAYS_FULL_RESET

        starts = np.array(
            [
                (trip.start - self.visa_start).days
                for trip in self.trips
            ],
            dtype=np.int16,
        )

        # For every reset index, for every trip, if the start index of the trip
        # is prior to the reset index, set all entries after the reset index to
        # 0.
        for reset in resets:
            addends[reset > starts, reset:] = 0

        days_left = START_DAYS + np.sum(addends, axis=0)
        return days_left

    @property
    def indexed_days_left(self) -> Dict[date, int]:
        return {
            when.astype('O'): left
            for when, left in zip(self.visa_dates, self.days_left)
        }

    @property
    def visa_dates(self) -> np.ndarray:
        return np.arange(self.visa_start, self.visa_end)

    @property
    def flat_lines(self) -> Iterable[str]:
        for when, days_left in zip(self.visa_dates, self.days_left):
            when_str = when.astype('O').strftime(DATE_FMT)
            yield f'{when_str} -- Days left: {days_left}'

    def print_flat(self) -> None:
        print('\n'.join(self.flat_lines))

    @property
    def padding_prefix(self) -> np.ndarray:
        month_start = self.visa_start.replace(day=1)
        return np.arange(month_start, self.visa_start)

    @property
    def padding_suffix(self) -> np.ndarray:
        if self.visa_end.day == 1:
            month_end = self.visa_end
        else:
            # The first day of the next month
            month_end = self.visa_end + relativedelta(months=1, day=1)

        return np.arange(self.visa_end, month_end)

    @property
    def padded_dates(self) -> Iterable[date]:
        for day in chain(self.padding_prefix, self.visa_dates, self.padding_suffix):
            yield day.astype('O')

    @property
    def padded_when(self) -> Iterable[str]:
        for day in self.padded_dates:
            yield f'{day.day:3d}'

    @property
    def padded_left(self) -> Iterable[str]:
        for _ in self.padding_prefix:
            yield '   '
        for left in self.days_left:
            yield f'{left:3d}'
        for _ in self.padding_suffix:
            yield '   '

    @staticmethod
    def print_week(
        day: date,
        cells: Iterable[str],
        first_week_epoch: Optional[date] = None,
    ) -> date:
        if first_week_epoch is None:
            row_day = day
        else:
            row_day = first_week_epoch

        if first_week_epoch is not None:
            while row_day < day:
                print('   ', end='')
                row_day += ONE_DAY

        for cell in cells:
            print(cell, end='')
            row_day += ONE_DAY
            if row_day.weekday() == SUNDAY or row_day.month != day.month:
                break
        print()

        return row_day

    def print_calendar(self) -> None:
        iso_days = (6, *range(6))
        week_header = ''.join(f'{day_abbr[d][:2]:>3}' for d in iso_days)

        whens = self.padded_when
        lefts = self.padded_left
        day, *_, end_day = self.padded_dates

        while day <= end_day:
            print(
                f'{day:%B %Y}\n'
                f'{week_header}'
            )
            # The Sunday before or on the first day of the given day's month
            month_start = day + relativedelta(day=1, weekday=SU(-1))

            self.print_week(day, whens, month_start)
            day = self.print_week(day, lefts, month_start)

            while True:
                self.print_week(day, whens)
                next_day = self.print_week(day, lefts)
                month_end = day.month != next_day.month
                day = next_day
                if month_end:
                    break
            print()


def test_op() -> None:
    db = TripDatabase(
        visa_start=date(2020, 11, 6),
        visa_end=date(2021, 11, 7),
        trips=(
            Trip(
                date(2020, 11, 8),
                # one day later than OP example due to being exclusive
                date(2020, 12, 29),
            ),
        ),
    )
    days = db.indexed_days_left

    assert date(2020, 11, 5) not in days
    assert days[date(2020, 11, 6)] == 90   # visa start
    assert days[date(2020, 11, 7)] == 90
    assert days[date(2020, 11, 8)] == 89   # trip start
    assert days[date(2020, 12, 27)] == 40
    assert days[date(2020, 12, 28)] == 39  # trip end
    assert days[date(2021, 3, 27)] == 39
    assert days[date(2021, 3, 28)] == 90   # reset
    assert days[date(2021, 11, 6)] == 90   # visa end
    assert date(2021, 11, 7) not in days


def test_180() -> None:
    db = TripDatabase(
        visa_start=date(2020, 1, 1),
        visa_end=date(2021, 1, 1),
        trips=(
            Trip(date(2020, 1, 3), date(2020, 2, 1)),  # 29 days
            Trip(date(2020, 3, 1), date(2020, 3, 8)),  # 7 days
            Trip(date(2020, 6, 1), date(2020, 6, 8)),  # 7 days
        ),
    )
    days = db.indexed_days_left

    # Visa start
    assert date(2019, 12, 31) not in days
    assert days[date(2020, 1, 1)] == 90

    # Trip A
    assert days[date(2020, 1, 2)] == 90
    assert days[date(2020, 1, 3)] == 89
    assert days[date(2020, 1, 30)] == 62
    assert days[date(2020, 1, 31)] == 61

    # In this idle period, there's neither enough time to reset nor to restore
    # individual days

    # Trip B
    assert days[date(2020, 2, 29)] == 61
    assert days[date(2020, 3, 1)] == 60
    assert days[date(2020, 3, 6)] == 55
    assert days[date(2020, 3, 7)] == 54

    # Idle period: same story

    # Trip C
    assert days[date(2020, 5, 29)] == 54
    assert days[date(2020, 6, 1)] == 53
    assert days[date(2020, 6, 6)] == 48
    assert days[date(2020, 6, 7)] == 47

    # Restoral from trip A
    assert days[date(2020, 7, 1)] == 47
    assert days[date(2020, 7, 2)] == 48
    assert days[date(2020, 7, 29)] == 75
    assert days[date(2020, 7, 30)] == 76

    # Restoral from trip B
    assert days[date(2020, 8, 28)] == 76
    assert days[date(2020, 8, 29)] == 77

    # In the middle of trip B restoral, reset occurs
    assert days[date(2020, 9, 4)] == 83
    assert days[date(2020, 9, 5)] == 90


if __name__ == '__main__':
    # db = TripDatabase.from_stdin() - if you want to enter this manually
    # db.print_flat() or print_calendar() for visual comparison
    test_op()
    test_180()

Calendar output (180 case)

January 2020
 Su Mo Tu We Th Fr Sa
           1  2  3  4
          90 90 89 88
  5  6  7  8  9 10 11
 87 86 85 84 83 82 81
 12 13 14 15 16 17 18
 80 79 78 77 76 75 74
 19 20 21 22 23 24 25
 73 72 71 70 69 68 67
 26 27 28 29 30 31
 66 65 64 63 62 61

February 2020
 Su Mo Tu We Th Fr Sa
                    1
                   61
  2  3  4  5  6  7  8
 61 61 61 61 61 61 61
  9 10 11 12 13 14 15
 61 61 61 61 61 61 61
 16 17 18 19 20 21 22
 61 61 61 61 61 61 61
 23 24 25 26 27 28 29
 61 61 61 61 61 61 61

March 2020
 Su Mo Tu We Th Fr Sa
  1  2  3  4  5  6  7
 60 59 58 57 56 55 54
  8  9 10 11 12 13 14
 54 54 54 54 54 54 54
 15 16 17 18 19 20 21
 54 54 54 54 54 54 54
 22 23 24 25 26 27 28
 54 54 54 54 54 54 54
 29 30 31
 54 54 54

April 2020
 Su Mo Tu We Th Fr Sa
           1  2  3  4
          54 54 54 54
  5  6  7  8  9 10 11
 54 54 54 54 54 54 54
 12 13 14 15 16 17 18
 54 54 54 54 54 54 54
 19 20 21 22 23 24 25
 54 54 54 54 54 54 54
 26 27 28 29 30
 54 54 54 54 54

May 2020
 Su Mo Tu We Th Fr Sa
                 1  2
                54 54
  3  4  5  6  7  8  9
 54 54 54 54 54 54 54
 10 11 12 13 14 15 16
 54 54 54 54 54 54 54
 17 18 19 20 21 22 23
 54 54 54 54 54 54 54
 24 25 26 27 28 29 30
 54 54 54 54 54 54 54
 31
 54

June 2020
 Su Mo Tu We Th Fr Sa
     1  2  3  4  5  6
    53 52 51 50 49 48
  7  8  9 10 11 12 13
 47 47 47 47 47 47 47
 14 15 16 17 18 19 20
 47 47 47 47 47 47 47
 21 22 23 24 25 26 27
 47 47 47 47 47 47 47
 28 29 30
 47 47 47

July 2020
 Su Mo Tu We Th Fr Sa
           1  2  3  4
          47 48 49 50
  5  6  7  8  9 10 11
 51 52 53 54 55 56 57
 12 13 14 15 16 17 18
 58 59 60 61 62 63 64
 19 20 21 22 23 24 25
 65 66 67 68 69 70 71
 26 27 28 29 30 31
 72 73 74 75 76 76

August 2020
 Su Mo Tu We Th Fr Sa
                    1
                   76
  2  3  4  5  6  7  8
 76 76 76 76 76 76 76
  9 10 11 12 13 14 15
 76 76 76 76 76 76 76
 16 17 18 19 20 21 22
 76 76 76 76 76 76 76
 23 24 25 26 27 28 29
 76 76 76 76 76 76 77
 30 31
 78 79

September 2020
 Su Mo Tu We Th Fr Sa
        1  2  3  4  5
       80 81 82 83 90
  6  7  8  9 10 11 12
 90 90 90 90 90 90 90
 13 14 15 16 17 18 19
 90 90 90 90 90 90 90
 20 21 22 23 24 25 26
 90 90 90 90 90 90 90
 27 28 29 30
 90 90 90 90

October 2020
 Su Mo Tu We Th Fr Sa
              1  2  3
             90 90 90
  4  5  6  7  8  9 10
 90 90 90 90 90 90 90
 11 12 13 14 15 16 17
 90 90 90 90 90 90 90
 18 19 20 21 22 23 24
 90 90 90 90 90 90 90
 25 26 27 28 29 30 31
 90 90 90 90 90 90 90

November 2020
 Su Mo Tu We Th Fr Sa
  1  2  3  4  5  6  7
 90 90 90 90 90 90 90
  8  9 10 11 12 13 14
 90 90 90 90 90 90 90
 15 16 17 18 19 20 21
 90 90 90 90 90 90 90
 22 23 24 25 26 27 28
 90 90 90 90 90 90 90
 29 30
 90 90

December 2020
 Su Mo Tu We Th Fr Sa
        1  2  3  4  5
       90 90 90 90 90
  6  7  8  9 10 11 12
 90 90 90 90 90 90 90
 13 14 15 16 17 18 19
 90 90 90 90 90 90 90
 20 21 22 23 24 25 26
 90 90 90 90 90 90 90
 27 28 29 30 31
 90 90 90 90 90

Answer 2

modular:

• use if __name__ == "__main__" guard to execute code, to allow imports
• use functions to encapsulate the logic and prevent unnecessary globals
• use class(es) to encapsulate the state -
• use constants as globals or class variables - the whole # Starting variables block
• consider naming the file __main__.py and adding a setup.py thus after installation you'll be able to do python -m <package name from setup.py> to start your program even without knowing the location

logical:

• trip_input.lower() == "end": can be replaced by just .strip() == "" and then you can do the common double-return/enter to end the input like for SMTP or likes and prevents relying on a specific, locale specific, string
• use standard format e.g. yyyy-mm-dd or likes that are defined for the whole world, such as ISO and you'll have less problems with parsing or converting to proper types - which will break e.g. when using a wide-character input methods (e.g. for Japanese) int(１) vs datetime.strptime("１", "%S") (%S Second as a decimal number [00,61])
• don't use indicies when not necessary, instead use unpacking because you might encounter a missing index thus IndexError
• move numbers such as 365 and strings such as "-" i.e. constants to constants - class/global variables because those are not changing which will allow you not to make a typo in the logic
• consider using dataclasses for the planned_trips.append(...) unless you can utilize datetime object directly with parsing from ISO format

visual:

• replace print() with logging
• use proper quotes where required - don't use \ as if you're in bash, use e.g. 'Hello "world"' or "Hello 'world'" strings instead
• use longer variables than i even for indicies, it'll actually help you later on if you know what that index actually represents; or if it's not an index then just use for _ in range(...) to "remove" the variable completely
• Ctrl-C will break your while True but also your program - consider using except KeyboardInterrupt or signal handling to exit from the program properly/cleanly and perhaps with a message
• nesting loops is ugly, consider moving it to functions that are named cleanly e.g. calculate_days_left() for the inner for?