How to type hint a `pl.date`?

July 8, 2024

Suppose we create some dates:

import polars as pl

df = pl.DataFrame(
    [
        pl.Series("start", ["2023-01-01"], dtype=pl.Date).str.to_date(),
        pl.Series("end", ["2024-01-01"], dtype=pl.Date).str.to_date(),
    ]
)

Now I can create a date range from these:

dates = pl.date_range(df[0, "start"], df[0, "end"], "1mo", eager=True)

But I want to define a function which takes a couple of dates and spits out a range, as a wrapper around pl.date_range:

def my_date_range(start: pl.Date, end: pl.Date) -> pl.Series:
    return pl.date_range(start, end, "1mo", eager=True)

The above doesn’t typecheck with pyright/Pylance, because:

Argument of type "Date" cannot be assigned to parameter "start" of type "IntoExprColumn | date | datetime" in function "date_range"
  Type "Date" is incompatible with type "IntoExprColumn | date | datetime"
    "Date" is incompatible with "date"
    "Date" is incompatible with "datetime"
    "Date" is incompatible with "Expr"
    "Date" is incompatible with "Series"
    "Date" is incompatible with "str"PylancereportArgumentType

If I check out type(df[0, "start"]), I see:

datetime.date

and pl.Date is no good because isinstance(df[0, "start"], pl.Date) == False.

I cannot figure out how to import datetime.date in order to use it as a type annotation (trying import polars.datetime as dt raises No module named 'polars.datetime').

How can this be done? Or put differently: how should my_date_range‘s date arguments be annotated?

>Solution :

Since datetime.date is compatible with the start and end parameters expected by pl.date_range() this should be sufficient:

import polars as pl
from datetime import date

def my_date_range(start: date, end: date) -> pl.Series:
    return pl.date_range(start, end, "1mo", eager=True)