Implement datetime timezone constraints

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)