Implement datetime timezone constraints (#343)

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.

Author: antonagestam

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',

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")]

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<DateTime>,
     gt: Option<DateTime>,
     now: Option<NowConstraint>,
+    tz: Option<TZConstraint>,
 }
 
 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<Self> {
+        match s {
+            "aware" => Ok(TZConstraint::Aware),
+            "naive" => Ok(TZConstraint::Naive),
+            _ => py_err!("Invalid tz_constraint {:?}", s),
+        }
+    }
+
+    pub fn from_py(schema: &PyDict) -> PyResult<Option<Self>> {
+        let py = schema.py();
+        match schema.get_as(intern!(py, "tz_constraint"))? {
+            Some(kind) => Ok(Some(Self::from_str(kind)?)),
+            None => Ok(None),
+        }
+    }
+}

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),

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)