13

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?

Martijn Pieters
  • 1,048,767
  • 296
  • 4,058
  • 3,343
Martin Thoma
  • 124,992
  • 159
  • 614
  • 958

1 Answers1

6

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.

Martin Thoma
  • 124,992
  • 159
  • 614
  • 958
  • 6
    You can avoid the cast by doing `NDT(datetime.datetime.now())`. NDT will be just the identity function at runtime, but at type check time acts as a constructor with a type signature of `def (val: datetime.datetime) -> NDT`. – Michael0x2a Apr 16 '20 at 01:11
  • Great! Even this keeps the correct type: `d1 + datetime.timedelta(1)` – pabouk - Ukraine stay strong May 31 '22 at 14:35