Support complex numbers

generate_self_schema.py

@@ -50,7 +50,7 @@
 def get_schema(obj: Any, definitions: dict[str, core_schema.CoreSchema]) -> core_schema.CoreSchema:  # noqa: C901
     if isinstance(obj, str):
         return {'type': obj}
-    elif obj in (datetime, timedelta, date, time, bool, int, float, str, decimal.Decimal):
+    elif obj in (datetime, timedelta, date, time, bool, int, float, str, decimal.Decimal, complex):
         return {'type': obj.__name__.lower()}
     elif is_typeddict(obj):
         return type_dict_schema(obj, definitions)

python/pydantic_core/core_schema.py

@@ -742,6 +742,44 @@ def decimal_schema(
     )
 
 
+class ComplexSchema(TypedDict, total=False):
+    type: Required[Literal['complex']]
+    ref: str
+    metadata: Any
+    serialization: SerSchema
+
+
+def complex_schema(
+    *,
+    ref: str | None = None,
+    metadata: Any = None,
+    serialization: SerSchema | None = None,
+) -> ComplexSchema:
+    """
+    Returns a schema that matches a complex value, e.g.:
+
+    ```py
+    from pydantic_core import SchemaValidator, core_schema
+
+    schema = core_schema.complex_schema()
+    v = SchemaValidator(schema)
+    assert v.validate_python('1+2j') == complex(1, 2)
+    assert v.validate_python(complex(1, 2)) == complex(1, 2)
+    ```
+
+    Args:
+        ref: optional unique identifier of the schema, used to reference the schema in other places
+        metadata: Any other information you want to include with the schema, not used by pydantic-core
+        serialization: Custom serialization schema
+    """
+    return _dict_not_none(
+        type='complex',
+        ref=ref,
+        metadata=metadata,
+        serialization=serialization,
+    )
+
+
 class StringSchema(TypedDict, total=False):
     type: Required[Literal['str']]
     pattern: Union[str, Pattern[str]]
@@ -3777,6 +3815,7 @@ def definition_reference_schema(
         DefinitionsSchema,
         DefinitionReferenceSchema,
         UuidSchema,
+        ComplexSchema,
     ]
 elif False:
     CoreSchema: TypeAlias = Mapping[str, Any]
@@ -3832,6 +3871,7 @@ def definition_reference_schema(
     'definitions',
     'definition-ref',
     'uuid',
+    'complex',
 ]
 
 CoreSchemaFieldType = Literal['model-field', 'dataclass-field', 'typed-dict-field', 'computed-field']
@@ -3936,6 +3976,8 @@ def definition_reference_schema(
     'decimal_max_digits',
     'decimal_max_places',
     'decimal_whole_digits',
+    'complex_parsing',
+    'complex_type',
 ]
 
 

src/errors/types.rs

@@ -422,6 +422,9 @@ error_types! {
     DecimalWholeDigits {
         whole_digits: {ctx_type: u64, ctx_fn: field_from_context},
     },
+    // Complex errors
+    ComplexParsing{},
+    ComplexType {},
 }
 
 macro_rules! render {
@@ -564,6 +567,8 @@ impl ErrorType {
             Self::DecimalMaxDigits {..} => "Decimal input should have no more than {max_digits} digit{expected_plural} in total",
             Self::DecimalMaxPlaces {..} => "Decimal input should have no more than {decimal_places} decimal place{expected_plural}",
             Self::DecimalWholeDigits {..} => "Decimal input should have no more than {whole_digits} digit{expected_plural} before the decimal point",
+            Self::ComplexParsing {..} => "Input should be a valid complex string following the rule at https://docs.python.org/3/library/functions.html#complex",
+            Self::ComplexType { .. } => "Input should be a valid complex number",
         }
     }
 

src/input/input_abstract.rs

@@ -9,7 +9,7 @@ use crate::lookup_key::{LookupKey, LookupPath};
 use crate::tools::py_err;
 
 use super::datetime::{EitherDate, EitherDateTime, EitherTime, EitherTimedelta};
-use super::return_enums::{EitherBytes, EitherInt, EitherString};
+use super::return_enums::{EitherBytes, EitherComplex, EitherInt, EitherString};
 use super::{EitherFloat, GenericIterator, ValidationMatch};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -172,6 +172,8 @@ pub trait Input<'py>: fmt::Debug + ToPyObject {
         strict: bool,
         microseconds_overflow_behavior: speedate::MicrosecondsPrecisionOverflowBehavior,
     ) -> ValMatch<EitherTimedelta<'py>>;
+
+    fn validate_complex(&self, py: Python<'py>) -> ValMatch<EitherComplex<'py>>;
 }
 
 /// The problem to solve here is that iterating collections often returns owned

src/input/input_json.rs

@@ -8,7 +8,9 @@ use speedate::MicrosecondsPrecisionOverflowBehavior;
 use strum::EnumMessage;
 
 use crate::errors::{ErrorType, ErrorTypeDefaults, InputValue, LocItem, ValError, ValResult};
+use crate::input::return_enums::EitherComplex;
 use crate::lookup_key::{LookupKey, LookupPath};
+use crate::validators::complex::string_to_complex;
 use crate::validators::decimal::create_decimal;
 
 use super::datetime::{
@@ -296,6 +298,18 @@ impl<'py, 'data> Input<'py> for JsonValue<'data> {
             _ => Err(ValError::new(ErrorTypeDefaults::TimeDeltaType, self)),
         }
     }
+
+    fn validate_complex(&self, py: Python<'py>) -> ValResult<ValidationMatch<EitherComplex<'py>>> {
+        match self {
+            JsonValue::Str(s) => Ok(ValidationMatch::strict(EitherComplex::Py(string_to_complex(
+                &PyString::new_bound(py, s),
+                self,
+            )?))),
+            JsonValue::Float(f) => Ok(ValidationMatch::strict(EitherComplex::Complex([*f, 0.0]))),
+            JsonValue::Int(f) => Ok(ValidationMatch::strict(EitherComplex::Complex([(*f) as f64, 0.0]))),
+            _ => Err(ValError::new(ErrorTypeDefaults::ComplexParsing, self)),
+        }
+    }
 }
 
 /// Required for JSON Object keys so the string can behave like an Input
@@ -425,6 +439,10 @@ impl<'py> Input<'py> for str {
     ) -> ValResult<ValidationMatch<EitherTimedelta<'py>>> {
         bytes_as_timedelta(self, self.as_bytes(), microseconds_overflow_behavior).map(ValidationMatch::lax)
     }
+
+    fn validate_complex(&self, _py: Python<'py>) -> ValResult<ValidationMatch<EitherComplex<'py>>> {
+        Err(ValError::new(ErrorTypeDefaults::ComplexType, self))
+    }
 }
 
 impl BorrowInput<'_> for &'_ String {

src/input/input_python.rs

@@ -6,15 +6,16 @@ use pyo3::prelude::*;
 
 use pyo3::types::PyType;
 use pyo3::types::{
-    PyBool, PyByteArray, PyBytes, PyDate, PyDateTime, PyDict, PyFloat, PyFrozenSet, PyInt, PyIterator, PyList,
-    PyMapping, PySet, PyString, PyTime, PyTuple,
+    PyBool, PyByteArray, PyBytes, PyComplex, PyDate, PyDateTime, PyDict, PyFloat, PyFrozenSet, PyInt, PyIterator,
+    PyList, PyMapping, PySet, PyString, PyTime, PyTuple,
 };
 
 use pyo3::PyTypeCheck;
 use speedate::MicrosecondsPrecisionOverflowBehavior;
 
 use crate::errors::{ErrorType, ErrorTypeDefaults, InputValue, LocItem, ValError, ValResult};
 use crate::tools::{extract_i64, safe_repr};
+use crate::validators::complex::string_to_complex;
 use crate::validators::decimal::{create_decimal, get_decimal_type};
 use crate::validators::Exactness;
 use crate::ArgsKwargs;
@@ -25,6 +26,7 @@ use super::datetime::{
     EitherTime,
 };
 use super::input_abstract::ValMatch;
+use super::return_enums::EitherComplex;
 use super::return_enums::{iterate_attributes, iterate_mapping_items, ValidationMatch};
 use super::shared::{
     decimal_as_int, float_as_int, get_enum_meta_object, int_as_bool, str_as_bool, str_as_float, str_as_int,
@@ -592,6 +594,25 @@ impl<'py> Input<'py> for Bound<'py, PyAny> {
 
         Err(ValError::new(ErrorTypeDefaults::TimeDeltaType, self))
     }
+
+    fn validate_complex<'a>(&'a self, _py: Python<'py>) -> ValResult<ValidationMatch<EitherComplex<'py>>> {
+        if let Ok(complex) = self.downcast::<PyComplex>() {
+            return Ok(ValidationMatch::exact(EitherComplex::Py(complex.to_owned())));
+        } else if let Ok(s) = self.downcast::<PyString>() {
+            return Ok(ValidationMatch::exact(EitherComplex::Py(string_to_complex(s, self)?)));
+        } else if self.is_exact_instance_of::<PyFloat>() {
+            return Ok(ValidationMatch::exact(EitherComplex::Complex([
+                self.extract::<f64>().unwrap(),
+                0.0,
+            ])));
+        } else if self.is_exact_instance_of::<PyInt>() {
+            return Ok(ValidationMatch::exact(EitherComplex::Complex([
+                self.extract::<i64>().unwrap() as f64,
+                0.0,
+            ])));
+        }
+        Err(ValError::new(ErrorTypeDefaults::ComplexType, self))
+    }
 }
 
 impl<'py> BorrowInput<'py> for Bound<'py, PyAny> {

src/input/input_string.rs

@@ -7,12 +7,14 @@ use crate::errors::{ErrorTypeDefaults, InputValue, LocItem, ValError, ValResult}
 use crate::input::py_string_str;
 use crate::lookup_key::{LookupKey, LookupPath};
 use crate::tools::safe_repr;
+use crate::validators::complex::string_to_complex;
 use crate::validators::decimal::create_decimal;
 
 use super::datetime::{
     bytes_as_date, bytes_as_datetime, bytes_as_time, bytes_as_timedelta, EitherDate, EitherDateTime, EitherTime,
 };
 use super::input_abstract::{Never, ValMatch};
+use super::return_enums::EitherComplex;
 use super::shared::{str_as_bool, str_as_float, str_as_int};
 use super::{
     Arguments, BorrowInput, EitherBytes, EitherFloat, EitherInt, EitherString, EitherTimedelta, GenericIterator, Input,
@@ -217,6 +219,13 @@ impl<'py> Input<'py> for StringMapping<'py> {
             Self::Mapping(_) => Err(ValError::new(ErrorTypeDefaults::TimeDeltaType, self)),
         }
     }
+
+    fn validate_complex(&self, _py: Python<'py>) -> ValResult<ValidationMatch<EitherComplex<'py>>> {
+        match self {
+            Self::String(s) => Ok(ValidationMatch::strict(EitherComplex::Py(string_to_complex(s, self)?))),
+            Self::Mapping(_) => Err(ValError::new(ErrorTypeDefaults::ComplexType, self)),
+        }
+    }
 }
 
 impl<'py> BorrowInput<'py> for StringMapping<'py> {

src/input/return_enums.rs

@@ -12,7 +12,7 @@ use pyo3::intern;
 use pyo3::prelude::*;
 #[cfg(not(PyPy))]
 use pyo3::types::PyFunction;
-use pyo3::types::{PyBytes, PyFloat, PyFrozenSet, PyIterator, PyMapping, PySet, PyString};
+use pyo3::types::{PyBytes, PyComplex, PyFloat, PyFrozenSet, PyIterator, PyMapping, PySet, PyString};
 
 use serde::{ser::Error, Serialize, Serializer};
 
@@ -716,3 +716,30 @@ impl ToPyObject for Int {
         }
     }
 }
+
+#[derive(Clone)]
+pub enum EitherComplex<'a> {
+    Complex([f64; 2]),
+    Py(Bound<'a, PyComplex>),
+}
+
+impl<'a> IntoPy<PyObject> for EitherComplex<'a> {
+    fn into_py(self, py: Python<'_>) -> PyObject {
+        match self {
+            Self::Complex(c) => PyComplex::from_doubles_bound(py, c[0], c[1]).into_py(py),
+            Self::Py(c) => c.into_py(py),
+        }
+    }
+}
+
+impl<'a> EitherComplex<'a> {
+    pub fn as_f64(&self, py: Python<'_>) -> [f64; 2] {
+        match self {
+            EitherComplex::Complex(f) => *f,
+            EitherComplex::Py(f) => [
+                f.getattr(intern!(py, "real")).unwrap().extract().unwrap(),
+                f.getattr(intern!(py, "imag")).unwrap().extract().unwrap(),
+            ],
+        }
+    }
+}

src/serializers/infer.rs

@@ -4,6 +4,7 @@ use pyo3::exceptions::PyTypeError;
 use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedStr;
+use pyo3::types::PyComplex;
 use pyo3::types::{PyByteArray, PyBytes, PyDict, PyFrozenSet, PyIterator, PyList, PySet, PyString, PyTuple};
 
 use serde::ser::{Error, Serialize, SerializeMap, SerializeSeq, Serializer};
@@ -226,6 +227,13 @@ pub(crate) fn infer_to_python_known(
                 }
                 PyList::new_bound(py, items).into_py(py)
             }
+            ObType::Complex => {
+                let dict = value.downcast::<PyDict>()?;
+                let new_dict = PyDict::new_bound(py);
+                let _ = new_dict.set_item("real", dict.get_item("real")?);
+                let _ = new_dict.set_item("imag", dict.get_item("imag")?);
+                new_dict.into_py(py)
+            }
             ObType::Path => value.str()?.into_py(py),
             ObType::Pattern => value.getattr(intern!(py, "pattern"))?.into_py(py),
             ObType::Unknown => {
@@ -274,6 +282,13 @@ pub(crate) fn infer_to_python_known(
                 );
                 iter.into_py(py)
             }
+            ObType::Complex => {
+                let dict = value.downcast::<PyDict>()?;
+                let new_dict = PyDict::new_bound(py);
+                let _ = new_dict.set_item("real", dict.get_item("real")?);
+                let _ = new_dict.set_item("imag", dict.get_item("imag")?);
+                new_dict.into_py(py)
+            }
             ObType::Unknown => {
                 if let Some(fallback) = extra.fallback {
                     let next_value = fallback.call1((value,))?;
@@ -402,6 +417,13 @@ pub(crate) fn infer_serialize_known<S: Serializer>(
         ObType::None => serializer.serialize_none(),
         ObType::Int | ObType::IntSubclass => serialize!(Int),
         ObType::Bool => serialize!(bool),
+        ObType::Complex => {
+            let v = value.downcast::<PyComplex>().map_err(py_err_se_err)?;
+            let mut map = serializer.serialize_map(Some(2))?;
+            map.serialize_entry(&"real", &v.real())?;
+            map.serialize_entry(&"imag", &v.imag())?;
+            map.end()
+        }
         ObType::Float | ObType::FloatSubclass => {
             let v = value.extract::<f64>().map_err(py_err_se_err)?;
             type_serializers::float::serialize_f64(v, serializer, extra.config.inf_nan_mode)
@@ -647,7 +669,7 @@ pub(crate) fn infer_json_key_known<'a>(
             }
             Ok(Cow::Owned(key_build.finish()))
         }
-        ObType::List | ObType::Set | ObType::Frozenset | ObType::Dict | ObType::Generator => {
+        ObType::List | ObType::Set | ObType::Frozenset | ObType::Dict | ObType::Generator | ObType::Complex => {
             py_err!(PyTypeError; "`{}` not valid as object key", ob_type)
         }
         ObType::Dataclass | ObType::PydanticSerializable => {

src/serializers/ob_type.rs

@@ -1,8 +1,8 @@
 use pyo3::prelude::*;
 use pyo3::sync::GILOnceCell;
 use pyo3::types::{
-    PyBool, PyByteArray, PyBytes, PyDate, PyDateTime, PyDelta, PyDict, PyFloat, PyFrozenSet, PyInt, PyIterator, PyList,
-    PyNone, PySet, PyString, PyTime, PyTuple, PyType,
+    PyBool, PyByteArray, PyBytes, PyComplex, PyDate, PyDateTime, PyDelta, PyDict, PyFloat, PyFrozenSet, PyInt,
+    PyIterator, PyList, PyNone, PySet, PyString, PyTime, PyTuple, PyType,
 };
 use pyo3::{intern, PyTypeInfo};
 
@@ -48,6 +48,7 @@ pub struct ObTypeLookup {
     pattern_object: PyObject,
     // uuid type
     uuid_object: PyObject,
+    complex: usize,
 }
 
 static TYPE_LOOKUP: GILOnceCell<ObTypeLookup> = GILOnceCell::new();
@@ -101,6 +102,7 @@ impl ObTypeLookup {
                 .to_object(py),
             pattern_object: py.import_bound("re").unwrap().getattr("Pattern").unwrap().to_object(py),
             uuid_object: py.import_bound("uuid").unwrap().getattr("UUID").unwrap().to_object(py),
+            complex: PyComplex::type_object_raw(py) as usize,
         }
     }
 
@@ -171,6 +173,7 @@ impl ObTypeLookup {
             ObType::Pattern => self.path_object.as_ptr() as usize == ob_type,
             ObType::Uuid => self.uuid_object.as_ptr() as usize == ob_type,
             ObType::Unknown => false,
+            ObType::Complex => self.complex == ob_type,
         };
 
         if ans {
@@ -425,6 +428,7 @@ pub enum ObType {
     Uuid,
     // unknown type
     Unknown,
+    Complex,
 }
 
 impl PartialEq for ObType {

src/serializers/shared.rs

@@ -142,6 +142,7 @@ combined_serializer! {
         Enum: super::type_serializers::enum_::EnumSerializer;
         Recursive: super::type_serializers::definitions::DefinitionRefSerializer;
         Tuple: super::type_serializers::tuple::TupleSerializer;
+        Complex: super::type_serializers::complex::ComplexSerializer;
     }
 }
 
@@ -251,6 +252,7 @@ impl PyGcTraverse for CombinedSerializer {
             CombinedSerializer::Recursive(inner) => inner.py_gc_traverse(visit),
             CombinedSerializer::Tuple(inner) => inner.py_gc_traverse(visit),
             CombinedSerializer::Uuid(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Complex(inner) => inner.py_gc_traverse(visit),
         }
     }
 }