Repository: pola-rs/polars
File: crates/polars-plan/src/dsl/functions/temporal.rs (Rust)

Code before:
    ///     .with_nanoseconds(nanoseconds)
    /// ```
    pub fn with_fractional_seconds(
        self,
        milliseconds: Expr,
        microseconds: Expr,
        nanoseconds: Expr,
    ) -> Self {
        Self {
            milliseconds,
            microseconds,
            nanoseconds,
            ..self
        }
    }

    impl_unit_setter!(with_weeks(weeks));
    impl_unit_setter!(with_days(days));
    impl_unit_setter!(with_hours(hours));
    impl_unit_setter!(with_minutes(minutes));
    impl_unit_setter!(with_seconds(seconds));
    impl_unit_setter!(with_milliseconds(milliseconds));
    impl_unit_setter!(with_microseconds(microseconds));
    impl_unit_setter!(with_nanoseconds(nanoseconds));

    fn as_literal(&self) -> Option<Expr> {
        use time_unit::convert_time_units;

        let extract_i64 = |e: &Expr| {
            let Expr::Literal(lv) = e else { return None };
            let av = lv.to_any_value()?;
            if !av.is_integer() {
                return None;
            };
            av.extract::<i64>()
        };
        let extract_f64 = |e: &Expr| {
            let Expr::Literal(lv) = e else { return None };
            let av = lv.to_any_value()?;
            if !av.is_float() {
                return None;
            };
            av.extract::<f64>()
        };

        let mut acc_i64 = Some(0);
        let mut try_add_to_i64 = |rhs| {
            acc_i64 = match (acc_i64, rhs) {
                (Some(acc), Some(x)) => Some(acc + x),
                _ => None,
            }

Review comment (bug):
If I read the code right, mixing integer and float literals will now cause `as_literal` to fail? I don't think that's intended?

Code after:
    ///     .with_nanoseconds(nanoseconds)
    /// ```
    pub fn with_fractional_seconds(
        self,
        milliseconds: Expr,
        microseconds: Expr,
        nanoseconds: Expr,
    ) -> Self {
        Self {
            milliseconds,
            microseconds,
            nanoseconds,
            ..self
        }
    }

    impl_unit_setter!(with_weeks(weeks));
    impl_unit_setter!(with_days(days));
    impl_unit_setter!(with_hours(hours));
    impl_unit_setter!(with_minutes(minutes));
    impl_unit_setter!(with_seconds(seconds));
    impl_unit_setter!(with_milliseconds(milliseconds));
    impl_unit_setter!(with_microseconds(microseconds));
    impl_unit_setter!(with_nanoseconds(nanoseconds));

    fn as_literal(&self) -> Option<Expr> {
        use time_unit::convert_time_units;

        let extract_i64 = |e: &Expr| {
            let Expr::Literal(lv) = e else { return None };
            let av = lv.to_any_value()?;
            if !av.is_integer() {
                return None;
            };
            av.extract::<i64>()
        };
        let extract_f64 = |e: &Expr| {
            let Expr::Literal(lv) = e else { return None };
            let av = lv.to_any_value()?;
            av.extract::<f64>()
        };

        let mut acc_i64 = Some(0);
        let mut try_add_to_i64 = |rhs| {
            acc_i64 = match (acc_i64, rhs) {
                (Some(acc), Some(x)) => Some(acc + x),
                _ => None,
            }
        };

        try_add_to_i64(extract_i64(&self.weeks).map(|v| v * 7 * NANOSECONDS_IN_DAY));