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)