From 9dccdea68fa8768d72cd9f678ed70488fc1dfbb1 Mon Sep 17 00:00:00 2001 From: Anton Agestam Date: Thu, 1 Dec 2022 23:25:59 +0100 Subject: [PATCH] Implement datetime timezone constraints This implements support for constraining datetime objects based on them having or not having timezone info. The aware kind constrains to objects that have timezone info, and symmetrically, the naive kind constrains to objects that do not have timezone info. --- pydantic_core/core_schema.py | 5 +++++ src/errors/types.rs | 4 ++++ src/validators/datetime.rs | 34 ++++++++++++++++++++++++++++++- tests/test_errors.py | 2 ++ tests/validators/test_datetime.py | 33 ++++++++++++++++++++++++++++++ 5 files changed, 77 insertions(+), 1 deletion(-) diff --git a/pydantic_core/core_schema.py b/pydantic_core/core_schema.py index 08768fb90..908360e35 100644 --- a/pydantic_core/core_schema.py +++ b/pydantic_core/core_schema.py @@ -276,6 +276,7 @@ class DatetimeSchema(TypedDict, total=False): lt: datetime gt: datetime now_op: Literal['past', 'future'] + tz_constraint: Literal['aware', 'naive'] # defaults to current local utc offset from `time.localtime().tm_gmtoff` # value is restricted to -86_400 < offset < 86_400 by bounds in generate_self_schema.py now_utc_offset: int @@ -291,6 +292,7 @@ def datetime_schema( lt: datetime | None = None, gt: datetime | None = None, now_op: Literal['past', 'future'] | None = None, + tz_constraint: Literal['aware', 'naive'] | None = None, now_utc_offset: int | None = None, ref: str | None = None, extra: Any = None, @@ -303,6 +305,7 @@ def datetime_schema( lt=lt, gt=gt, now_op=now_op, + tz_constraint=tz_constraint, now_utc_offset=now_utc_offset, ref=ref, extra=extra, @@ -1226,6 +1229,8 @@ def multi_host_url_schema( 'datetime_object_invalid', 'datetime_past', 'datetime_future', + 'datetime_aware', + 'datetime_naive', 'time_delta_type', 'time_delta_parsing', 'frozen_set_type', diff --git a/src/errors/types.rs b/src/errors/types.rs index 77f5cbf29..2f80e8186 100644 --- a/src/errors/types.rs +++ b/src/errors/types.rs @@ -261,6 +261,10 @@ pub enum ErrorType { DatetimePast, #[strum(message = "Datetime should be in the future")] DatetimeFuture, + #[strum(message = "Datetime should have timezone info")] + DatetimeAware, + #[strum(message = "Datetime should not have timezone info")] + DatetimeNaive, // --------------------- // timedelta errors #[strum(message = "Input should be a valid timedelta")] diff --git a/src/validators/datetime.rs b/src/validators/datetime.rs index e5f2fca65..d775f10e3 100644 --- a/src/validators/datetime.rs +++ b/src/validators/datetime.rs @@ -92,6 +92,12 @@ impl Validator for DateTimeValidator { } } } + + match (&constraints.tz, speedate_dt.offset) { + (Some(TZConstraint::Aware), None) => return Err(ValError::new(ErrorType::DatetimeAware, input)), + (Some(TZConstraint::Naive), Some(_)) => return Err(ValError::new(ErrorType::DatetimeNaive, input)), + _ => (), + } } Ok(datetime.try_into_py(py)?) } @@ -108,6 +114,7 @@ struct DateTimeConstraints { ge: Option, gt: Option, now: Option, + tz: Option, } impl DateTimeConstraints { @@ -119,8 +126,9 @@ impl DateTimeConstraints { ge: py_datetime_as_datetime(schema, intern!(py, "ge"))?, gt: py_datetime_as_datetime(schema, intern!(py, "gt"))?, now: NowConstraint::from_py(schema)?, + tz: TZConstraint::from_py(schema)?, }; - if c.le.is_some() || c.lt.is_some() || c.ge.is_some() || c.gt.is_some() || c.now.is_some() { + if c.le.is_some() || c.lt.is_some() || c.ge.is_some() || c.gt.is_some() || c.now.is_some() || c.tz.is_some() { Ok(Some(c)) } else { Ok(None) @@ -199,3 +207,27 @@ impl NowConstraint { } } } + +#[derive(Debug, Clone)] +pub enum TZConstraint { + Aware, + Naive, +} + +impl TZConstraint { + pub fn from_str(s: &str) -> PyResult { + match s { + "aware" => Ok(TZConstraint::Aware), + "naive" => Ok(TZConstraint::Naive), + _ => py_err!("Invalid tz_constraint {:?}", s), + } + } + + pub fn from_py(schema: &PyDict) -> PyResult> { + let py = schema.py(); + match schema.get_as(intern!(py, "tz_constraint"))? { + Some(kind) => Ok(Some(Self::from_str(kind)?)), + None => Ok(None), + } + } +} diff --git a/tests/test_errors.py b/tests/test_errors.py index 807d252fd..0b4a3f3a0 100644 --- a/tests/test_errors.py +++ b/tests/test_errors.py @@ -235,6 +235,8 @@ def f(input_value, **kwargs): ('datetime_object_invalid', 'Invalid datetime object, got foobar', {'error': 'foobar'}), ('datetime_past', 'Datetime should be in the past', None), ('datetime_future', 'Datetime should be in the future', None), + ('datetime_aware', 'Datetime should have timezone info', None), + ('datetime_naive', 'Datetime should not have timezone info', None), ('time_delta_type', 'Input should be a valid timedelta', None), ('time_delta_parsing', 'Input should be a valid timedelta, foobar', {'error': 'foobar'}), ('frozen_set_type', 'Input should be a valid frozenset', None), diff --git a/tests/validators/test_datetime.py b/tests/validators/test_datetime.py index 82fd01939..4b872caf5 100644 --- a/tests/validators/test_datetime.py +++ b/tests/validators/test_datetime.py @@ -366,3 +366,36 @@ def test_mock_utc_offset_8_hours(mocker): def test_offset_too_large(): with pytest.raises(SchemaError, match=r'Input should be greater than -86400 \[type=greater_than,'): SchemaValidator(core_schema.datetime_schema(now_op='past', now_utc_offset=-24 * 3600)) + + +class TestTZConstraints: + aware_validator = SchemaValidator(core_schema.datetime_schema(tz_constraint='aware')) + naive_validator = SchemaValidator(core_schema.datetime_schema(tz_constraint='naive')) + + def test_raises_schema_error_for_unknown_constraint_kind(self): + with pytest.raises( + SchemaError, + match=( + r'Input should be \'aware\' or \'naive\' ' + r'\[type=literal_error, input_value=\'foo\', input_type=str\]' + ), + ): + SchemaValidator({'type': 'datetime', 'tz_constraint': 'foo'}) + + def test_can_validate_aware_value(self): + value = datetime.now(tz=timezone.utc) + assert value is self.aware_validator.validate_python(value) + + def test_raises_validation_error_when_aware_given_naive(self): + value = datetime.now() + with pytest.raises(ValidationError, match=r'Datetime should have timezone info'): + assert self.aware_validator.validate_python(value) + + def test_can_validate_naive_value(self): + value = datetime.now() + assert value is self.naive_validator.validate_python(value) + + def test_raises_validation_error_when_naive_given_aware(self): + value = datetime.now(tz=timezone.utc) + with pytest.raises(ValidationError, match=r'Datetime should not have timezone info'): + assert self.naive_validator.validate_python(value)