DateTime - Protocol possible with current typing features?

[randolf-scholz]

So, I was trying to write a Protocol that captures the essence of datetime-like data types. (In particular datetime.datetime, numpy.datetime64 and pandas.Timestamp). The problem is how to annotate the __add__ and __sub__-methods. Each datetime-type has a paired timedelta type (datetime.timedelta, numpy.timedelta64, pandas.Timedelta), and it's unclear to me how to deal with that. This is what I have so far, it passes pyright with warnings (TD is unbound), but mypy hates it.

from typing import (
    Protocol,
    Self,
    SupportsFloat,
    SupportsInt,
    TypeVar,
    overload,
    reveal_type,
)

class TimeDelta(Protocol):
    """Time delta provides several arithmetical operations."""

    # unary operations
    def __pos__(self) -> Self: ...
    def __neg__(self) -> Self: ...
    def __abs__(self) -> Self: ...

    # comparisons
    def __eq__(self, other: object, /) -> bool: ...
    def __ne__(self, other: object, /) -> bool: ...
    def __le__(self, other: Self, /) -> bool: ...
    def __lt__(self, other: Self, /) -> bool: ...
    def __ge__(self, other: Self, /) -> bool: ...
    def __gt__(self, other: Self, /) -> bool: ...

    # arithmetic
    # addition +
    def __add__(self, other: Self, /) -> Self: ...
    def __radd__(self, other: Self, /) -> Self: ...

    # subtraction -
    def __sub__(self, other: Self, /) -> Self: ...
    def __rsub__(self, other: Self, /) -> Self: ...

    # multiplication *
    def __mul__(self, other: int, /) -> Self: ...
    def __rmul__(self, other: int, /) -> Self: ...

    # division /
    def __truediv__(self, other: Self, /) -> SupportsFloat: ...

    # @overload
    # def __truediv__(self, other: Self, /) -> float: ...
    # @overload
    # def __truediv__(self, other: float, /) -> Self: ...

    # floor division //
    def __floordiv__(self, other: Self, /) -> SupportsInt: ...

    # @overload
    # def __floordiv__(self, other: Self, /) -> int: ...
    # @overload
    # def __floordiv__(self, other: int, /) -> Self: ...

    # modulo %
    def __mod__(self, other: Self, /) -> Self: ...

    # NOTE: __rmod__ missing on fallback pydatetime
    # def __rmod__(self, other: Self, /) -> Self: ...

    # divmod
    def __divmod__(self, other: Self, /) -> tuple[SupportsInt, Self]: ...

    # NOTE: __rdivmod__ missing on fallback pydatetime
    # def __rdivmod__(self, other: Self, /) -> tuple[SupportsInt, Self]: ...


TD = TypeVar("TD", bound=TimeDelta)


class DateTime(Protocol):  # bind appropriate TimeDelta type
    """Datetime can be compared and subtracted."""

    def __hash__(self) -> int: ...

    def __eq__(self, other: object, /) -> bool: ...

    def __le__(self, other: Self, /) -> bool: ...
    def __lt__(self, other: Self, /) -> bool: ...
    def __ge__(self, other: Self, /) -> bool: ...
    def __gt__(self, other: Self, /) -> bool: ...

    def __add__(self, other: TD, /) -> Self: ...
    def __radd__(self, other: TD, /) -> Self: ...

    @overload
    def __sub__(self, other: TD, /) -> Self: ...
    @overload
    def __sub__(self, other: Self, /) -> TD: ...


def main():
    from datetime import datetime, timedelta

    import numpy
    import pandas

    ISO_DATE = "2021-01-01"

    TIMEDELTAS: dict[str, TimeDelta] = {
        "float": 10.0,
        "int": 10,
        "numpy": numpy.timedelta64(1, "D"),
        "numpy_float": numpy.float64(10.0),
        "numpy_int": numpy.int64(10),
        "pandas": pandas.Timedelta(days=1),
        "python": timedelta(days=1),
        # "arrow": pyarrow.scalar(TD, type=pyarrow_td_type),
    }

    DATETIMES: dict[str, DateTime] = {
        "float": 10.0,
        "int": 10,
        "numpy": numpy.datetime64(ISO_DATE),
        "numpy_float": numpy.float64(10.0),
        "numpy_int": numpy.int64(10),
        "pandas": pandas.Timestamp(ISO_DATE),
        "python": datetime.fromisoformat(ISO_DATE),
        # "arrow": pyarrow.scalar(DT, type=pyarrow.timestamp("s")),
    }

    x: DateTime = 1.0
    y: DateTime = 10
    z: DateTime = datetime.fromisoformat(ISO_DATE)
    w: DateTime = numpy.datetime64(ISO_DATE)
    f: DateTime = numpy.float64(10.0)
    d = datetime.fromisoformat(ISO_DATE)
    e: DateTime = datetime.fromisoformat(ISO_DATE)
    reveal_type(d - d)
    reveal_type(e - e)
    reveal_type(d + (d - d))
    reveal_type(e + (e - e))
    reveal_type(w - w)

    def foo(x: DateTime):
        reveal_type(x - x)

[Daverball]

I'm not sure this makes a whole lot of sense, since this would allow mixing incompatible DateTime and TimeDelta types. There's no proper way to express how these two types are linked to one another in a protocol (without making it generic), so mypy is actually doing the right thing and rejecting these objects as DateTime, since their __add__/__radd__/__sub__ require a specific type, rather than a Protocol, it would not be type safe to accept them.

If you want to make this work despite this, then you have to use Any instead of TD in the annotation of these methods. Alternatively you could make the Protocol generic in TD, then you could use DateTime[Any] to get this behavior, albeit slightly more type safe than actually using Any in the protocol.

If I change DateTime(Protocol) to DateTime(Protocol[TD]) mypy stops complaining.

[Daverball]

Maybe an example to illustrate in what way these protocols are harmful:

a: DateTime = datetime(2020, 1, 1)
b: TimeDelta = timedelta(days=5)
x = a + 5.0
y = b + 5.0

Both of these would be a perfectly legal operation according to these protocols, but will result in a runtime exception.

[randolf-scholz]

Oh wow, that is indeed a complete showstopper, then. Does that mean that even something like a Protocol for vectors isn't really feasible? Doesn't this run into the exact same problem:

from typing import Protocol, Self, TypeVar

scalar = TypeVar("scalar", bound=float)

class Vector(Protocol[scalar]):
    """Represents an element of a vector space."""

    def __len__(self) -> int: ...
    def __getitem__(self, key: int, /) -> scalar: ...
    def __pos__(self) -> Self: ...
    def __neg__(self) -> Self: ...
    def __add__(self, other: Self, /) -> Self: ...
    def __sub__(self, other: Self, /) -> Self: ...
    def __mul__(self, other: scalar, /) -> Self: ...

[Daverball]

It depends somewhat on the implementation and what the contract of the API is. But generally this pattern doesn't work, yeah, since a structural type forgets the concrete type it was matched against. There are some fringe use-cases where a protocol like this might be useful (such as an upper bound on a TypeVar), but then SupportsLenAndGetitem[int, scalar] would probably be just as useful and avoids allowing dangerous operations in the implementation without explicit type narrowing.

[randolf-scholz]

The motivation is, I have some quite generic code like for instance this function:

def compute_grid(tmin, tmax, timedelta, t0=None) -> list[int]:
    r"""Compute $\{k∈ℤ ∣ tₘᵢₙ ≤ t₀+k⋅Δt ≤ tₘₐₓ\}$.""" 
    if timedelta == (tmin - tmin): # zero-variable of the correct type.
        raise ValueError("Δt=0 is not allowed!")

    t0 = tmin if t0 is None else t0
    kmax = math.floor((tmax - t0) / timedelta)
    kmin = math.ceil((tmin - t0) / timedelta)
    return list(range(kmin, kmax + 1))

Which should work for any compatible datetime/timedelta variables. Of course, I could add overloads for (datetime.datetime, datetime.timedelta), (numpy.datetime64, numpy.timedelta64) and (pandas.Timestamp, pandas.Timedelta), and even (float, float), but that seems to really go against the spirit of duck-typing. So I was hoping for the possibility of annotating this with a protocol.

I have used this function for all 4 of these in the past, it's a helper function used for some messy data science.

[Daverball]

This is inherently difficult to model, because for type safety all the tX parameters should have the same type, so you would like to use a TypeVar for this, but the return type of tX.__sub__(tY) needs to match the type of timedelta. This would require higher kinded TypeVars that can bind to another type var, i.e. you would need to be able to do something like:

TD = TypeVar('TD', bound=TimeDelta)
D = TypeVar('D', bound=DateTime[TD])
def compute_grid(tmin: D[TD], tmax: D[TD], timedelta: TD, t0: D[TD] | None = None) -> list[int]:

So either you have to give up on type safety or use overloads for all the valid type combinations you want to support.