What is the recommended way to annotate datetime objects?

Question

Suppose I have a function which takes two datetimes and returns the difference in seconds:

import datetime


def diff(d1: datetime.datetime, d2: datetime.datetime) -> float:
    return (d2 - d1).total_seconds()


if __name__ == '__main__':
    d1 = datetime.datetime.now()
    d2 = datetime.datetime.now(datetime.timezone.utc)
    print(diff(d1, d2))

mypy tells me this is fine:

$ python3.8 -m mypy test.py
Success: no issues found in 1 source file

But I get a TypeError:

TypeError: can't subtract offset-naive and offset-aware datetimes

The reason is stated clearly in the error message. The type annotation was not good enough.

I guess this is a pretty common case. Is there a recommended way to annotate timezone aware vs unaware datetime objects which makes proper use of mypy?

Answer 1

One potential solution is using TypeVar, NewType and cast:

import datetime
from typing import NewType, TypeVar, cast

NDT = NewType("NDT", datetime.datetime)  # non-aware datetime
ADT = NewType("ADT", datetime.datetime)  # timezone aware datetime
DatetimeLike = TypeVar("DatetimeLike", NDT, ADT)


def diff(d1: DatetimeLike, d2: DatetimeLike) -> float:
    return (d2 - d1).total_seconds()


if __name__ == "__main__":
    d1: NDT = cast(NDT, datetime.datetime.now())
    d2: ADT = cast(ADT, datetime.datetime.now(datetime.timezone.utc))

    # Fails with:
    #    error: Value of type variable "DatetimeLike" of "diff" cannot be "datetime"
    # You have to use either NDT or ADT
    print(diff(d1, d2))

I don't like about this solution that you have to use cast and that the error message is less clear than it was before.