add timedelta type (#137)

* add timedelta type

* fix lint single quotes

* fix lint single quotes

* fix review

* merge main branch

* add timedelta testcase

* fix timedelta positive

* fix pytimedelta_as_timedelta

* remove dbg

* use approx comparison for total_seconds

* uprev speedate to limit durations

Co-authored-by: Samuel Colvin <[email protected]>

Author: fcfangcc

Cargo.lock

@@ -17,7 +17,7 @@ enum_dispatch = "0.3.8"
 serde = "1.0.137"
 indexmap = "1.8.1"
 mimalloc = { version = "0.1.29", default-features = false, optional = true }
-speedate = "0.4.1"
+speedate = "0.6.0"
 ahash = "0.7.6"
 nohash-hasher = "0.2.0"
 

Cargo.toml

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import sys
-from datetime import date, datetime, time
+from datetime import date, datetime, time, timedelta
 from typing import Any, Callable, Dict, List, Type, Union
 
 if sys.version_info < (3, 11):
@@ -211,6 +211,16 @@ class DatetimeSchema(TypedDict, total=False):
     ref: str
 
 
+class TimedeltaSchema(TypedDict, total=False):
+    type: Required[Literal['timedelta']]
+    strict: bool
+    le: timedelta
+    ge: timedelta
+    lt: timedelta
+    gt: timedelta
+    ref: str
+
+
 class TupleFixLenSchema(TypedDict, total=False):
     type: Required[Literal['tuple-fix-len']]
     items_schema: Required[List[Schema]]
@@ -287,6 +297,7 @@ class CallableSchema(TypedDict):
     DateSchema,
     TimeSchema,
     DatetimeSchema,
+    TimedeltaSchema,
     IsInstanceSchema,
     CallableSchema,
 ]

pydantic_core/_types.py

@@ -190,6 +190,12 @@ pub enum ErrorKind {
     )]
     DateTimeObjectInvalid { error: String },
     // ---------------------
+    // timedelta errors
+    #[strum(message = "Value must be a valid timedelta")]
+    TimeDeltaType,
+    #[strum(message = "Value must be a valid timedelta, {error}")]
+    TimeDeltaParsing { error: &'static str },
+    // ---------------------
     // frozenset errors
     #[strum(message = "Value must be a valid frozenset")]
     FrozenSetType,
@@ -275,6 +281,7 @@ impl ErrorKind {
             Self::TimeParsing { error } => render!(template, error),
             Self::DateTimeParsing { error } => render!(template, error),
             Self::DateTimeObjectInvalid { error } => render!(template, error),
+            Self::TimeDeltaParsing { error } => render!(template, error),
             Self::IsInstanceOf { class } => render!(template, class),
             _ => template.to_string(),
         }
@@ -321,6 +328,7 @@ impl ErrorKind {
             Self::TimeParsing { error } => py_dict!(py, error),
             Self::DateTimeParsing { error } => py_dict!(py, error),
             Self::DateTimeObjectInvalid { error } => py_dict!(py, error),
+            Self::TimeDeltaParsing { error } => py_dict!(py, error),
             Self::IsInstanceOf { class } => py_dict!(py, class),
             _ => Ok(None),
         }

src/errors/kinds.rs

@@ -1,12 +1,13 @@
 use pyo3::intern;
 use pyo3::prelude::*;
-use pyo3::types::{PyDate, PyDateTime, PyDelta, PyTime, PyTzInfo};
-use speedate::{Date, DateTime, Time};
+use pyo3::types::{PyDate, PyDateTime, PyDelta, PyDeltaAccess, PyTime, PyTzInfo};
+use speedate::{Date, DateTime, Duration, Time};
 use strum::EnumMessage;
 
-use super::Input;
 use crate::errors::{ErrorKind, ValError, ValResult};
 
+use super::Input;
+
 pub enum EitherDate<'a> {
     Raw(Date),
     Py(&'a PyDate),
@@ -69,6 +70,73 @@ impl<'a> From<&'a PyTime> for EitherTime<'a> {
     }
 }
 
+#[derive(Debug, Clone)]
+pub enum EitherTimedelta<'a> {
+    Raw(Duration),
+    Py(&'a PyDelta),
+}
+
+impl<'a> From<Duration> for EitherTimedelta<'a> {
+    fn from(timedelta: Duration) -> Self {
+        Self::Raw(timedelta)
+    }
+}
+
+impl<'a> From<&'a PyDelta> for EitherTimedelta<'a> {
+    fn from(timedelta: &'a PyDelta) -> Self {
+        Self::Py(timedelta)
+    }
+}
+
+pub fn pytimedelta_as_timedelta(py_timedelta: &PyDelta) -> Duration {
+    // see https://docs.python.org/3/c-api/datetime.html#c.PyDateTime_DELTA_GET_DAYS
+    // days can be negative, but seconds and microseconds are always positive.
+    let mut days = py_timedelta.get_days(); // -999999999 to 999999999
+    let mut seconds = py_timedelta.get_seconds(); // 0 through 86399
+    let mut microseconds = py_timedelta.get_microseconds(); // 0 through 999999
+    let positive = days >= 0;
+    if !positive {
+        // negative timedelta, we need to adjust values to match duration logic
+        if microseconds != 0 {
+            seconds += 1;
+            microseconds = (microseconds - 1_000_000).abs();
+        }
+        if seconds != 0 {
+            days += 1;
+            seconds = (seconds - 86400).abs();
+        }
+        days = days.abs();
+    }
+    // we can safely "unwrap" since the methods above guarantee values are in the correct ranges.
+    Duration::new(positive, days as u32, seconds as u32, microseconds as u32).unwrap()
+}
+
+impl<'a> EitherTimedelta<'a> {
+    pub fn as_raw(&self) -> Duration {
+        match self {
+            Self::Raw(timedelta) => timedelta.clone(),
+            Self::Py(py_timedelta) => pytimedelta_as_timedelta(py_timedelta),
+        }
+    }
+
+    pub fn try_into_py(self, py: Python<'_>) -> PyResult<PyObject> {
+        let timedelta = match self {
+            Self::Py(timedelta) => Ok(timedelta),
+            Self::Raw(duration) => {
+                let sign = if duration.positive { 1 } else { -1 };
+                PyDelta::new(
+                    py,
+                    sign * duration.day as i32,
+                    sign * duration.second as i32,
+                    sign * duration.microsecond as i32,
+                    true,
+                )
+            }
+        }?;
+        Ok(timedelta.into_py(py))
+    }
+}
+
 macro_rules! pytime_as_time {
     ($py_time:expr) => {
         speedate::Time {
@@ -287,6 +355,36 @@ pub fn float_as_time<'a>(input: &'a impl Input<'a>, timestamp: f64) -> ValResult
     int_as_time(input, timestamp.floor() as i64, microseconds.round() as u32)
 }
 
+pub fn bytes_as_timedelta<'a, 'b>(input: &'a impl Input<'a>, bytes: &'b [u8]) -> ValResult<'a, EitherTimedelta<'a>> {
+    match Duration::parse_bytes(bytes) {
+        Ok(dt) => Ok(dt.into()),
+        Err(err) => Err(ValError::new(
+            ErrorKind::TimeDeltaParsing {
+                error: err.get_documentation().unwrap_or_default(),
+            },
+            input,
+        )),
+    }
+}
+
+pub fn int_as_duration(total_seconds: i64) -> Duration {
+    let positive = total_seconds >= 0;
+    let total_seconds = total_seconds.unsigned_abs();
+    // we can safely unwrap here since we've guaranteed seconds and microseconds can't cause overflow
+    let days = (total_seconds / 86400) as u32;
+    let seconds = (total_seconds % 86400) as u32;
+    Duration::new(positive, days, seconds, 0).unwrap()
+}
+
+pub fn float_as_duration(total_seconds: f64) -> Duration {
+    let positive = total_seconds >= 0_f64;
+    let total_seconds = total_seconds.abs();
+    let microsecond = total_seconds.fract() * 1_000_000.0;
+    let days = (total_seconds / 86400f64) as u32;
+    let seconds = total_seconds as u64 % 86400;
+    Duration::new(positive, days, seconds as u32, microsecond.round() as u32).unwrap()
+}
+
 #[pyclass(module = "pydantic_core._pydantic_core", extends = PyTzInfo)]
 #[derive(Debug, Clone)]
 struct TzInfo {

src/input/datetime.rs

@@ -6,7 +6,7 @@ use pyo3::types::PyType;
 use crate::errors::{InputValue, LocItem, ValResult};
 use crate::input::datetime::EitherTime;
 
-use super::datetime::{EitherDate, EitherDateTime};
+use super::datetime::{EitherDate, EitherDateTime, EitherTimedelta};
 use super::return_enums::{EitherBytes, EitherString};
 use super::{GenericMapping, GenericSequence};
 
@@ -99,6 +99,12 @@ pub trait Input<'a>: fmt::Debug + ToPyObject {
         self.strict_tuple()
     }
 
+    fn strict_timedelta(&self) -> ValResult<EitherTimedelta>;
+
+    fn lax_timedelta(&self) -> ValResult<EitherTimedelta> {
+        self.strict_timedelta()
+    }
+
     fn is_instance(&self, _class: &PyType) -> PyResult<bool> {
         Ok(false)
     }

src/input/input_abstract.rs

@@ -3,11 +3,11 @@ use pyo3::types::PyType;
 use crate::errors::{ErrorKind, InputValue, LocItem, ValError, ValResult};
 
 use super::datetime::{
-    bytes_as_date, bytes_as_datetime, bytes_as_time, float_as_datetime, float_as_time, int_as_datetime, int_as_time,
-    EitherDate, EitherDateTime, EitherTime,
+    bytes_as_date, bytes_as_datetime, bytes_as_time, bytes_as_timedelta, float_as_datetime, float_as_duration,
+    float_as_time, int_as_datetime, int_as_duration, int_as_time, EitherDate, EitherDateTime, EitherTime,
 };
 use super::shared::{float_as_int, int_as_bool, str_as_bool, str_as_int};
-use super::{EitherBytes, EitherString, GenericMapping, GenericSequence, Input, JsonInput};
+use super::{EitherBytes, EitherString, EitherTimedelta, GenericMapping, GenericSequence, Input, JsonInput};
 
 impl<'a> Input<'a> for JsonInput {
     /// This is required by since JSON object keys are always strings, I don't think it can be called
@@ -197,6 +197,22 @@ impl<'a> Input<'a> for JsonInput {
             _ => Err(ValError::new(ErrorKind::TupleType, self)),
         }
     }
+
+    fn strict_timedelta(&self) -> ValResult<EitherTimedelta> {
+        match self {
+            JsonInput::String(v) => bytes_as_timedelta(self, v.as_bytes()),
+            _ => Err(ValError::new(ErrorKind::TimeDeltaType, self)),
+        }
+    }
+
+    fn lax_timedelta(&self) -> ValResult<EitherTimedelta> {
+        match self {
+            JsonInput::String(v) => bytes_as_timedelta(self, v.as_bytes()),
+            JsonInput::Int(v) => Ok(int_as_duration(*v).into()),
+            JsonInput::Float(v) => Ok(float_as_duration(*v).into()),
+            _ => Err(ValError::new(ErrorKind::TimeDeltaType, self)),
+        }
+    }
 }
 
 /// Required for Dict keys so the string can behave like an Input
@@ -298,4 +314,8 @@ impl<'a> Input<'a> for String {
     fn strict_tuple<'data>(&'data self) -> ValResult<GenericSequence<'data>> {
         Err(ValError::new(ErrorKind::TupleType, self))
     }
+
+    fn strict_timedelta(&self) -> ValResult<super::EitherTimedelta> {
+        bytes_as_timedelta(self, self.as_bytes())
+    }
 }

src/input/input_json.rs

@@ -3,19 +3,20 @@ use std::str::from_utf8;
 use pyo3::exceptions::{PyAttributeError, PyTypeError};
 use pyo3::prelude::*;
 use pyo3::types::{
-    PyBool, PyByteArray, PyBytes, PyDate, PyDateTime, PyDict, PyFrozenSet, PyInt, PyList, PyMapping, PySequence, PySet,
-    PyString, PyTime, PyTuple, PyType,
+    PyBool, PyByteArray, PyBytes, PyDate, PyDateTime, PyDelta, PyDict, PyFrozenSet, PyInt, PyList, PyMapping,
+    PySequence, PySet, PyString, PyTime, PyTuple, PyType,
 };
 use pyo3::{intern, AsPyPointer};
 
 use crate::errors::{as_internal, py_err_string, ErrorKind, InputValue, LocItem, ValError, ValResult};
 
 use super::datetime::{
-    bytes_as_date, bytes_as_datetime, bytes_as_time, date_as_datetime, float_as_datetime, float_as_time,
-    int_as_datetime, int_as_time, EitherDate, EitherDateTime, EitherTime,
+    bytes_as_date, bytes_as_datetime, bytes_as_time, bytes_as_timedelta, date_as_datetime, float_as_datetime,
+    float_as_duration, float_as_time, int_as_datetime, int_as_duration, int_as_time, EitherDate, EitherDateTime,
+    EitherTime,
 };
 use super::shared::{float_as_int, int_as_bool, str_as_bool, str_as_int};
-use super::{repr_string, EitherBytes, EitherString, GenericMapping, GenericSequence, Input};
+use super::{repr_string, EitherBytes, EitherString, EitherTimedelta, GenericMapping, GenericSequence, Input};
 
 impl<'a> Input<'a> for PyAny {
     fn as_loc_item(&'a self) -> LocItem {
@@ -391,6 +392,30 @@ impl<'a> Input<'a> for PyAny {
         }
     }
 
+    fn strict_timedelta(&self) -> ValResult<EitherTimedelta> {
+        if let Ok(dt) = self.cast_as::<PyDelta>() {
+            Ok(dt.into())
+        } else {
+            Err(ValError::new(ErrorKind::TimeDeltaType, self))
+        }
+    }
+
+    fn lax_timedelta(&self) -> ValResult<EitherTimedelta> {
+        if let Ok(dt) = self.cast_as::<PyDelta>() {
+            Ok(dt.into())
+        } else if let Ok(py_str) = self.cast_as::<PyString>() {
+            bytes_as_timedelta(self, py_str.to_string_lossy().as_bytes())
+        } else if let Ok(py_bytes) = self.cast_as::<PyBytes>() {
+            bytes_as_timedelta(self, py_bytes.as_bytes())
+        } else if let Ok(int) = self.extract::<i64>() {
+            Ok(int_as_duration(int).into())
+        } else if let Ok(float) = self.extract::<f64>() {
+            Ok(float_as_duration(float).into())
+        } else {
+            Err(ValError::new(ErrorKind::TimeDeltaType, self))
+        }
+    }
+
     fn is_instance(&self, class: &PyType) -> PyResult<bool> {
         self.is_instance(class)
     }

src/input/input_python.rs

@@ -8,7 +8,7 @@ mod parse_json;
 mod return_enums;
 mod shared;
 
-pub use datetime::{EitherDate, EitherDateTime, EitherTime};
+pub use datetime::{EitherDate, EitherDateTime, EitherTime, EitherTimedelta};
 pub use input_abstract::Input;
 pub use parse_json::{JsonInput, JsonObject};
 pub use return_enums::{EitherBytes, EitherString, GenericMapping, GenericSequence};

src/input/mod.rs

@@ -32,10 +32,10 @@ mod recursive;
 mod set;
 mod string;
 mod time;
+mod timedelta;
 mod tuple;
 mod typed_dict;
 mod union;
-
 #[pyclass(module = "pydantic_core._pydantic_core")]
 #[derive(Debug, Clone)]
 pub struct SchemaValidator {
@@ -292,6 +292,8 @@ pub fn build_validator<'a>(
         datetime::DateTimeValidator,
         // frozensets
         frozenset::FrozenSetValidator,
+        // timedelta
+        timedelta::TimeDeltaValidator,
         // introspection types
         is_instance::IsInstanceValidator,
         callable::CallableValidator,
@@ -374,6 +376,8 @@ pub enum CombinedValidator {
     Datetime(datetime::DateTimeValidator),
     // frozensets
     FrozenSet(frozenset::FrozenSetValidator),
+    // timedelta
+    Timedelta(timedelta::TimeDeltaValidator),
     // introspection types
     IsInstance(is_instance::IsInstanceValidator),
     Callable(callable::CallableValidator),