Support parsing timezone-aware datetimes in constructors when data type is also timezone-aware

[stinodego]

Description

I ran into this today, and I think we can improve behavior here.

Consider this code:

from datetime import datetime
from zoneinfo import ZoneInfo
import polars as pl

data = [datetime(1970, 1, 1, tzinfo=ZoneInfo("Iran"))]
s = pl.Series(data, dtype=pl.Datetime(time_zone="Iran"), strict=False)

ValueError: time-zone-aware datetimes are converted to UTC

Please either drop the time zone from the dtype, or set it to 'UTC'. To convert to a different time zone, please use `.dt.convert_time_zone`.

This is odd. If we're casting timezone-aware data anyway, might as well cast it to desired time zone, right?

One of the benefits of doing this is that timezone-aware data can then be roundtripped, like in one of our tests:

@given(s=series())
def test_to_list(s: pl.Series) -> None:
    values = s.to_list()
    result = pl.Series(values, dtype=s.dtype)
    assert_series_equal(s, result, categorical_as_str=True)

For reference, PyArrow seems to handle this a bit differently from us and they do support this:

import pyarrow as pa

data = [datetime(1970, 1, 1, tzinfo=ZoneInfo("Iran"))]
arr = pa.array(data, type=pa.timestamp("us", tz="Iran"))

print(arr)  # 1969-12-31 20:30:00.000000Z
print(arr.type)  # timestamp[us, tz=Iran]

I may be missing something here, but I thought I'd throw this out there. Let's see what @MarcoGorelli thinks 😄

[MarcoGorelli]

thanks for the ping

so, for .str.to_datetime, the current rules are:

• naive string, time zone => parse into the given time zone

  In [17]: pl.Series(['2020-01-01T01:23:45']).str.to_datetime(time_zone='Asia/Kathmandu')
  Out[17]:
  shape: (1,)
  Series: '' [datetime[μs, Asia/Kathmandu]]
  [
        2020-01-01 01:23:45 +0545
  ]

• offset-aware string, no time zone set: convert to UTC:

  
  In [18]: pl.Series(['2020-01-01T01:23:45+05:45']).str.to_datetime()
  Out[18]:
  shape: (1,)
  Series: '' [datetime[μs, UTC]]
  [
        2019-12-31 19:38:45 UTC
  ]

• naive, no time zone: parse as naive
• offset-aware string, time zone set: raise

  ComputeError: offset-aware datetimes are converted to UTC. Please either drop the time zone from the function call, or set it to UTC. To convert to a different time zone, please use `convert_time_zone`.

The constructor should probably not be too different. Let's see:

• naive datetime, time zone => parse into given time zone. ✅ same

  In [20]: pl.Series([datetime(2020, 1, 1)], dtype=pl.Datetime('us', 'Asia/Kathmandu'))
  Out[20]:
  shape: (1,)
  Series: '' [datetime[μs, Asia/Kathmandu]]
  [
        2020-01-01 00:00:00 +0545
  ]

• time zone aware datetime, no time zone given => convert to UTC. ✅ same

  In [21]: pl.Series([datetime(2020, 1, 1, tzinfo=ZoneInfo('Asia/Kathmandu'))], dtype=pl.Datetime)
  <ipython-input-21-3f19ea488109>:1: TimeZoneAwareConstructorWarning: Constructing a Series with time-zone-aware datetimes results in a Series with UTC time zone. To silence this warning, you can filter warnings of class TimeZoneAwareConstructorWarning, or set 'UTC' as the time zone of your datatype.
  pl.Series([datetime(2020, 1, 1, tzinfo=ZoneInfo('Asia/Kathmandu'))], dtype=pl.Datetime)
  Out[21]:
  shape: (1,)
  Series: '' [datetime[μs, UTC]]
  [
        2019-12-31 18:15:00 UTC
  ]

• naive, no time zone: parse as naive. ✅ same
• time-zone-aware, time zone set: also raises (as you found), so it's also ✅ same

For the last one, I think you're suggesting to convert to the given time zone. So long as it's clearly documented, and it's done for both the Series constructor and .str.to_datetime, I think it makes sense