polars_arrow/legacy/kernels/
time.rs

1use std::str::FromStr;
2
3#[cfg(feature = "timezones")]
4use chrono::{LocalResult, NaiveDateTime, TimeZone};
5#[cfg(feature = "timezones")]
6use chrono_tz::Tz;
7#[cfg(feature = "timezones")]
8use polars_error::PolarsResult;
9use polars_error::{polars_bail, PolarsError};
10#[cfg(feature = "serde")]
11use serde::{Deserialize, Serialize};
12use strum_macros::IntoStaticStr;
13
14pub enum Ambiguous {
15    Earliest,
16    Latest,
17    Null,
18    Raise,
19}
20impl FromStr for Ambiguous {
21    type Err = PolarsError;
22
23    fn from_str(s: &str) -> Result<Self, Self::Err> {
24        match s {
25            "earliest" => Ok(Ambiguous::Earliest),
26            "latest" => Ok(Ambiguous::Latest),
27            "raise" => Ok(Ambiguous::Raise),
28            "null" => Ok(Ambiguous::Null),
29            s => polars_bail!(InvalidOperation:
30                "Invalid argument {}, expected one of: \"earliest\", \"latest\", \"null\", \"raise\"", s
31            ),
32        }
33    }
34}
35
36#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, IntoStaticStr)]
37#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
38#[strum(serialize_all = "snake_case")]
39pub enum NonExistent {
40    Null,
41    Raise,
42}
43
44#[cfg(feature = "timezones")]
45pub fn convert_to_naive_local(
46    from_tz: &Tz,
47    to_tz: &Tz,
48    ndt: NaiveDateTime,
49    ambiguous: Ambiguous,
50    non_existent: NonExistent,
51) -> PolarsResult<Option<NaiveDateTime>> {
52    let ndt = from_tz.from_utc_datetime(&ndt).naive_local();
53    match to_tz.from_local_datetime(&ndt) {
54        LocalResult::Single(dt) => Ok(Some(dt.naive_utc())),
55        LocalResult::Ambiguous(dt_earliest, dt_latest) => match ambiguous {
56            Ambiguous::Earliest => Ok(Some(dt_earliest.naive_utc())),
57            Ambiguous::Latest => Ok(Some(dt_latest.naive_utc())),
58            Ambiguous::Null => Ok(None),
59            Ambiguous::Raise => {
60                polars_bail!(ComputeError: "datetime '{}' is ambiguous in time zone '{}'. Please use `ambiguous` to tell how it should be localized.", ndt, to_tz)
61            },
62        },
63        LocalResult::None => match non_existent {
64            NonExistent::Raise => polars_bail!(ComputeError:
65                "datetime '{}' is non-existent in time zone '{}'. You may be able to use `non_existent='null'` to return `null` in this case.",
66                ndt, to_tz
67            ),
68            NonExistent::Null => Ok(None),
69        },
70    }
71}
72
73/// Same as convert_to_naive_local, but return `None` instead
74/// raising - in some cases this can be used to save a string allocation.
75#[cfg(feature = "timezones")]
76pub fn convert_to_naive_local_opt(
77    from_tz: &Tz,
78    to_tz: &Tz,
79    ndt: NaiveDateTime,
80    ambiguous: Ambiguous,
81) -> Option<Option<NaiveDateTime>> {
82    let ndt = from_tz.from_utc_datetime(&ndt).naive_local();
83    match to_tz.from_local_datetime(&ndt) {
84        LocalResult::Single(dt) => Some(Some(dt.naive_utc())),
85        LocalResult::Ambiguous(dt_earliest, dt_latest) => match ambiguous {
86            Ambiguous::Earliest => Some(Some(dt_earliest.naive_utc())),
87            Ambiguous::Latest => Some(Some(dt_latest.naive_utc())),
88            Ambiguous::Null => Some(None),
89            Ambiguous::Raise => None,
90        },
91        LocalResult::None => None,
92    }
93}