Add schema validation to PyDict -> Document

Author: GodTamIt

src/document.rs

@@ -13,7 +13,7 @@ use chrono::{offset::TimeZone, Utc};
 
 use tantivy as tv;
 
-use crate::{facet::Facet, to_pyerr};
+use crate::{facet::Facet, schema::Schema, to_pyerr};
 use serde_json::Value as JsonValue;
 use std::{
     collections::{BTreeMap, HashMap},
@@ -200,6 +200,71 @@ pub(crate) fn extract_value(any: &PyAny) -> PyResult<Value> {
     Err(to_pyerr(format!("Value unsupported {any:?}")))
 }
 
+pub(crate) fn extract_value_for_type(
+    any: &PyAny,
+    tv_type: tv::schema::Type,
+    field_name: &str,
+) -> PyResult<Value> {
+    // Helper function to create `PyErr`s returned by this function.
+    fn to_pyerr_for_type<'a, E: std::error::Error>(
+        type_name: &'a str,
+        field_name: &'a str,
+        any: &'a PyAny,
+    ) -> impl Fn(E) -> PyErr + 'a {
+        move |_| {
+            to_pyerr(format!(
+                "Expected {} type for field {}, got {:?}",
+                type_name, field_name, any
+            ))
+        }
+    }
+
+    let value = match tv_type {
+        tv::schema::Type::Str => Value::Str(
+            any.extract::<String>()
+                .map_err(to_pyerr_for_type("Str", field_name, any))?,
+        ),
+        tv::schema::Type::U64 => Value::U64(
+            any.extract::<u64>()
+                .map_err(to_pyerr_for_type("U64", field_name, any))?,
+        ),
+        tv::schema::Type::I64 => Value::I64(
+            any.extract::<i64>()
+                .map_err(to_pyerr_for_type("I64", field_name, any))?,
+        ),
+        tv::schema::Type::F64 => Value::F64(
+            any.extract::<f64>()
+                .map_err(to_pyerr_for_type("F64", field_name, any))?,
+        ),
+        tv::schema::Type::Date => {
+            let py_datetime = any
+                .downcast::<PyDateTime>()
+                .map_err(to_pyerr_for_type("DateTime", field_name, any))?;
+            let datetime = Utc
+                .with_ymd_and_hms(
+                    py_datetime.get_year(),
+                    py_datetime.get_month().into(),
+                    py_datetime.get_day().into(),
+                    py_datetime.get_hour().into(),
+                    py_datetime.get_minute().into(),
+                    py_datetime.get_second().into(),
+                )
+                .single()
+                .unwrap();
+
+            Value::Date(tv::DateTime::from_timestamp_secs(datetime.timestamp()))
+        }
+        tv::schema::Type::Facet => Value::Facet(
+            any.extract::<Facet>()
+                .map_err(to_pyerr_for_type("Facet", field_name, any))?
+                .inner,
+        ),
+        _ => return Err(to_pyerr(format!("Value unsupported {:?}", any))),
+    };
+
+    Ok(value)
+}
+
 fn extract_value_single_or_list(any: &PyAny) -> PyResult<Vec<Value>> {
     if let Ok(values) = any.downcast::<PyList>() {
         values.iter().map(extract_value).collect()
@@ -208,51 +273,135 @@ fn extract_value_single_or_list(any: &PyAny) -> PyResult<Vec<Value>> {
     }
 }
 
-#[pymethods]
-impl Document {
-    #[new]
-    #[pyo3(signature = (**kwargs))]
-    fn new(kwargs: Option<&PyDict>) -> PyResult<Self> {
-        let mut document = Document::default();
-        if let Some(field_dict) = kwargs {
-            document.extend(field_dict)?;
+fn extract_value_single_or_list_for_type(
+    any: &PyAny,
+    field_type: &tv::schema::FieldType,
+    field_name: &str,
+) -> PyResult<Vec<Value>> {
+    // Check if a numeric fast field supports multivalues.
+    let allow_multivalues = match field_type {
+        tv::schema::FieldType::U64(opts)
+        | tv::schema::FieldType::I64(opts)
+        | tv::schema::FieldType::F64(opts) => opts
+            .get_fastfield_cardinality()
+            .map(|card| card == tv::schema::Cardinality::MultiValues)
+            .unwrap_or(true),
+        _ => true,
+    };
+
+    if let Ok(values) = any.downcast::<PyList>() {
+        let result = values
+            .iter()
+            .map(|any| {
+                extract_value_for_type(any, field_type.value_type(), field_name)
+            })
+            .collect::<PyResult<Vec<Value>>>()?;
+
+        if !allow_multivalues && result.len() > 1 {
+            return Err(to_pyerr(format!(
+                "Field '{}' does not support multivalues",
+                field_name
+            )));
+        } else {
+            Ok(result)
         }
-        Ok(document)
+    } else {
+        Ok(vec![extract_value_for_type(
+            any,
+            field_type.value_type(),
+            field_name,
+        )?])
     }
+}
+
+impl Document {
+    fn extract_py_values_from_dict(
+        py_dict: &PyDict,
+        schema: Option<&Schema>,
+        out_field_values: &mut BTreeMap<String, Vec<tv::schema::Value>>,
+    ) -> PyResult<()> {
+        // TODO: Reserve when https://github.com/rust-lang/rust/issues/72631 is stable.
+        // out_field_values.reserve(py_dict.len());
 
-    fn extend(&mut self, py_dict: &PyDict) -> PyResult<()> {
-        let mut field_values: BTreeMap<String, Vec<tv::schema::Value>> =
-            BTreeMap::new();
         for key_value_any in py_dict.items() {
             if let Ok(key_value) = key_value_any.downcast::<PyTuple>() {
                 if key_value.len() != 2 {
                     continue;
                 }
-                let key: String = key_value.get_item(0)?.extract()?;
-                let value_list =
-                    extract_value_single_or_list(key_value.get_item(1)?)?;
-                field_values.insert(key, value_list);
+                let key = key_value.get_item(0)?.extract::<String>()?;
+
+                let field_type = if let Some(schema) = schema {
+                    let field_type =
+                        schema.inner.get_field(key.as_str()).map(|field| {
+                            schema.inner.get_field_entry(field).field_type()
+                        });
+
+                    if let Some(field_type) = field_type {
+                        // A field type was found, so validate it after the values are extracted.
+                        Some(field_type)
+                    } else {
+                        // The field does not exist in the schema, so skip over it.
+                        continue;
+                    }
+                } else {
+                    // No schema was provided, so do not validate anything.
+                    None
+                };
+
+                let value_list = if let Some(field_type) = field_type {
+                    extract_value_single_or_list_for_type(
+                        key_value.get_item(1)?,
+                        field_type,
+                        key.as_str(),
+                    )?
+                } else {
+                    extract_value_single_or_list(key_value.get_item(1)?)?
+                };
+
+                out_field_values.insert(key, value_list);
             }
         }
-        self.field_values.extend(field_values.into_iter());
+
         Ok(())
     }
+}
+
+#[pymethods]
+impl Document {
+    #[new]
+    #[pyo3(signature = (**kwargs))]
+    fn new(kwargs: Option<&PyDict>) -> PyResult<Self> {
+        let mut document = Document::default();
+        if let Some(field_dict) = kwargs {
+            document.extend(field_dict, None)?;
+        }
+        Ok(document)
+    }
+
+    fn extend(
+        &mut self,
+        py_dict: &PyDict,
+        schema: Option<&Schema>,
+    ) -> PyResult<()> {
+        Document::extract_py_values_from_dict(
+            py_dict,
+            schema,
+            &mut self.field_values,
+        )
+    }
 
     #[staticmethod]
-    fn from_dict(py_dict: &PyDict) -> PyResult<Document> {
+    fn from_dict(
+        py_dict: &PyDict,
+        schema: Option<&Schema>,
+    ) -> PyResult<Document> {
         let mut field_values: BTreeMap<String, Vec<tv::schema::Value>> =
             BTreeMap::new();
-        for key_value_any in py_dict.items() {
-            if let Ok(key_value) = key_value_any.downcast::<PyTuple>() {
-                if key_value.len() != 2 {
-                    continue;
-                }
-                let key: String = key_value.get_item(0)?.extract()?;
-                let value_list =
-                    extract_value_single_or_list(key_value.get_item(1)?)?;
-                field_values.insert(key, value_list);
-            }
-        }
+        Document::extract_py_values_from_dict(
+            py_dict,
+            schema,
+            &mut field_values,
+        )?;
         Ok(Document { field_values })
     }