Make TzInfo picklable (#770)

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

Author: adriangb

python/pydantic_core/_pydantic_core.pyi

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import datetime
 import decimal
 import sys
 from typing import Any, Callable, Generic, Optional, Type, TypeVar
@@ -43,6 +44,7 @@ __all__ = [
     'to_json',
     'to_jsonable_python',
     'list_all_errors',
+    'TzInfo',
 ]
 __version__: str
 build_profile: str
@@ -323,3 +325,10 @@ def list_all_errors() -> list[ErrorTypeInfo]:
     """
     Get information about all built-in errors.
     """
+
+@final
+class TzInfo(datetime.tzinfo):
+    def tzname(self, _dt: datetime.datetime | None) -> str | None: ...
+    def utcoffset(self, _dt: datetime.datetime | None) -> datetime.timedelta: ...
+    def dst(self, _dt: datetime.datetime | None) -> datetime.timedelta: ...
+    def __deepcopy__(self, _memo: dict[Any, Any]) -> 'TzInfo': ...

src/input/datetime.rs

@@ -1,9 +1,11 @@
 use pyo3::intern;
 use pyo3::prelude::*;
+
 use pyo3::types::{PyDate, PyDateTime, PyDelta, PyDeltaAccess, PyDict, PyTime, PyTzInfo};
 use speedate::MicrosecondsPrecisionOverflowBehavior;
 use speedate::{Date, DateTime, Duration, ParseError, Time, TimeConfig};
 use std::borrow::Cow;
+
 use strum::EnumMessage;
 
 use super::Input;
@@ -463,7 +465,7 @@ pub fn float_as_duration<'a>(input: &'a impl Input<'a>, total_seconds: f64) -> V
 #[pyclass(module = "pydantic_core._pydantic_core", extends = PyTzInfo)]
 #[derive(Clone)]
 #[cfg_attr(debug_assertions, derive(Debug))]
-struct TzInfo {
+pub struct TzInfo {
     seconds: i32,
 }
 
@@ -502,4 +504,10 @@ impl TzInfo {
     fn __deepcopy__(&self, py: Python, _memo: &PyDict) -> PyResult<Py<Self>> {
         Py::new(py, self.clone())
     }
+
+    pub fn __reduce__(&self, py: Python) -> PyResult<PyObject> {
+        let args = (self.seconds,);
+        let cls = Py::new(py, self.clone())?.getattr(py, "__class__")?;
+        Ok((cls, args).into_py(py))
+    }
 }

src/input/mod.rs

@@ -10,6 +10,7 @@ mod parse_json;
 mod return_enums;
 mod shared;
 
+pub use datetime::TzInfo;
 pub(crate) use datetime::{
     duration_as_pytimedelta, pydate_as_date, pydatetime_as_datetime, pytime_as_time, pytimedelta_as_duration,
     EitherDate, EitherDateTime, EitherTime, EitherTimedelta,

src/lib.rs

@@ -24,6 +24,7 @@ mod url;
 mod validators;
 
 // required for benchmarks
+pub use self::input::TzInfo;
 pub use self::url::{PyMultiHostUrl, PyUrl};
 pub use argument_markers::{ArgsKwargs, PydanticUndefinedType};
 pub use build_tools::SchemaError;
@@ -93,6 +94,7 @@ fn _pydantic_core(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<PyMultiHostUrl>()?;
     m.add_class::<ArgsKwargs>()?;
     m.add_class::<SchemaSerializer>()?;
+    m.add_class::<TzInfo>()?;
     m.add_function(wrap_pyfunction!(to_json, m)?)?;
     m.add_function(wrap_pyfunction!(to_jsonable_python, m)?)?;
     m.add_function(wrap_pyfunction!(list_all_errors, m)?)?;

tests/validators/test_datetime.py

@@ -1,9 +1,11 @@
 import copy
 import json
+import pickle
 import platform
 import re
 from datetime import date, datetime, time, timedelta, timezone, tzinfo
 from decimal import Decimal
+from typing import Dict
 
 import pytest
 import pytz
@@ -476,3 +478,38 @@ def test_tz_constraint_too_high():
 def test_tz_constraint_wrong():
     with pytest.raises(SchemaError, match="Input should be 'aware' or 'naive"):
         SchemaValidator(core_schema.datetime_schema(tz_constraint='wrong'))
+
+
+def test_tz_pickle() -> None:
+    """
+    https://github.com/pydantic/pydantic-core/issues/589
+    """
+    v = SchemaValidator(core_schema.datetime_schema())
+    original = datetime(2022, 6, 8, 12, 13, 14, tzinfo=timezone(timedelta(hours=-12, minutes=-15)))
+    validated = v.validate_python('2022-06-08T12:13:14-12:15')
+    assert validated == original
+    assert pickle.loads(pickle.dumps(validated)) == validated == original
+
+
+def test_tz_hash() -> None:
+    v = SchemaValidator(core_schema.datetime_schema())
+    lookup: Dict[datetime, str] = {}
+    for day in range(1, 10):
+        input_str = f'2022-06-{day:02}T12:13:14-12:15'
+        validated = v.validate_python(input_str)
+        lookup[validated] = input_str
+
+    assert len(lookup) == 9
+    assert (
+        lookup[datetime(2022, 6, 8, 12, 13, 14, tzinfo=timezone(timedelta(hours=-12, minutes=-15)))]
+        == '2022-06-08T12:13:14-12:15'
+    )
+
+
+def test_tz_cmp() -> None:
+    v = SchemaValidator(core_schema.datetime_schema())
+    validated1 = v.validate_python('2022-06-08T12:13:14-12:15')
+    validated2 = v.validate_python('2022-06-08T12:13:14-12:14')
+
+    assert validated1 > validated2
+    assert validated2 < validated1